PyTorch 2.0: 新特性与升级指南
什么是 PyTorch 2.0?
PyTorch 2.0 是 PyTorch 的最新版本，它保留了之前版本的即时执行模式（eager mode），同时引入了一个全新的编译模式。这个编译模式通过 torch.compile 函数实现，有潜力显著提升模型的训练和推理速度。
为什么是 2.0 而不是 1.14?
PyTorch 团队认为这个版本引入的新特性足以改变用户使用 PyTorch 的方式，因此决定将其命名为 2.0 而不是 1.14。
如何安装 PyTorch 2.0?
你可以通过 pip 安装最新的 nightly 版本。根据你的 CUDA 版本或是否使用 CPU，选择相应的安装命令：
# CUDA 11.8
pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu118
# CUDA 11.7
pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cu117
# CPU
pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --index-url https://download.pytorch.org/whl/nightly/cpu
2.0 版本的兼容性如何?
PyTorch 2.0 完全向后兼容 1.x 版本。你无需修改现有的 PyTorch 工作流程。只需添加一行代码 model = torch.compile(model) 就可以优化你的模型以使用 2.0 的新特性。
如何迁移到 PyTorch 2.0?
大多数情况下，你的代码无需任何改动就可以在 PyTorch 2.0 上运行。如果你想使用新的编译模式特性，只需要在你的模型上调用 torch.compile：
import torch
def train(model, dataloader):
model = torch.compile(model)
for batch in dataloader:
run_epoch(model, batch)
def infer(model, input):
model = torch.compile(model)
return model(**input)
PyTorch 2.0 的工作原理
当你使用 torch.compile(model) 包装你的模型时，模型会经历以下三个步骤：
- 图获取：模型被重写为子图块。
- 图降低：PyTorch 操作被分解为特定后端的核心操作。
- 图编译：核心操作调用相应的低级设备特定操作。
PyTorch 2.0 的新组件
- TorchDynamo：从 Python 字节码生成 FX 图。
- AOTAutograd：为 TorchDynamo 捕获的前向图生成对应的反向图。
- PrimTorch：将复杂的 PyTorch 操作分解为更简单和基本的操作。
- 后端：与 TorchDynamo 集成，将图编译为可在加速器上运行的 IR。
分布式训练
在编译模式下，DDP 和 FSDP 可以比即时执行模式快 15%（FP32）到 80%（AMP 精度）。使用 DDP 时，请确保设置 static_graph=False。
遇到问题怎么办？
如果你的代码在编译模式下运行变慢或崩溃，很可能是由于图断裂（graph breaks）导致的。你可以参考 PyTorch 官方文档 来诊断和解决这些问题。
PyTorch 2.0 带来了显著的性能提升和新特性，同时保持了与旧版本的兼容性。通过简单的一行代码，你就可以享受到这些优化带来的好处.