文章目录[隐藏]
关注程旭源，我们一起向上生长
本文是为了让大家快速入门和使用Torch 2.x，所以主要结合代码去介绍如何使用TorchDynamo和TorchInductor，并探索其原理。
Pytorch官方团队宣布Pyotch 2.X的新时代已经到来。
主要四个方面：
-
TorchDynamo：把Python Bytecode 字节码分析生成 FX 图；
-
TorchInductor：神经网络编译器；
-
AOTAutograd ：以 ahead-of-time 的方式生成反向图；
-
PrimTorch ：引入一个小型算子集，使后端更容易。
torch.compile 编译模式，在不更改模型代码的情况下加速模型。
TorchDynamo相比fx和jit的trace，可以更方便地跟踪data-dependent分支，而且不像torch.script一样需要修改代码。TorchInductor就是编译器，有对应的IR（define-by-run IR），然后可以利用triton生成高性能的算子。Triton是一个可以使用Python编写高性能GPU程序的Python库。
使用方式
使用函数 torch.compile 来包装模型，并返回一个编译后的模型。
函数解析：
-
mode: 指定编译器在编译时应该优化什么。预设defaul模式，也可以设置reduce-overhead或者max-autotune。
-
dynamic: 指定是否启用动态形状还是静态形状。
-
fullgraph：是否将整个程序编译成一个图形。
-
backend: 指定要使用的编译器后端。默认Torch Inductor。
使用示例
只需要红框中一行代码即可，其他和Pytorch1.x的相同。
在训练和推理中示例：
序列化
我们可以序列化optimized_model或原模型的状态字典。方法相同。但目前不能直接保存optimized_model。
推理和导出
我们可以通过使用一行装饰器torch._dynamo.optimize()，轻松地尝试不同的编译器后端，使PyTorch代码更快。
安装TorchDynamo
GPU：安装GPU PyTorch TorchDynamo以及GPU TorchDynamo依赖项(用于CUDA 11.7)
pip install numpy --pre torch[dynamo] --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117pip install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu
原理图：
Hugging Face 和 Pytotch 2.x
我们使用Bert作为示例，其实也就增加了两行代码。
https://github.com/pytorch/pytorch/blob/master/torch/_dynamo/optimizations/backends.py自定义后端backend
加速后端backend
多个编译器组合使用
组合多个一起玩也是可以的。玩得溜、玩的花。
【知乎、B站同名】
<