PyTorch2.0来了！100%向后兼容，一行代码将训练提速76%！

点击下方卡片&＃xff0c;关注“CVer”公众号

AI/CV重磅干货&＃xff0c;第一时间送达

点击进入—> CV 微信技术交流群

转载自&＃xff1a;机器之心

PyTorch 官方&＃xff1a;我们这次的新特性太好用了&＃xff0c;所以就直接叫 2.0 了。

前段时间&＃xff0c;PyTorch 团队在官方博客宣布 Pytorch 1.13 发布&＃xff0c;包含 BetterTransformer 稳定版等多项更新。在体验新特性的同时&＃xff0c;不少人也在期待下一个版本的推出。

出乎意料的是&＃xff0c;这个新版本很快就来了&＃xff0c;而且是跨越式的 2.0 版&＃xff01;

新版本的重要进步体现在速度和可用性&＃xff0c;而且完全向后兼容。PyTorch 团队表示&＃xff0c;PyTorch 2.0 是他们向 2.x 系列迈出的第一步&＃xff0c;其稳定版预计在 2023 年 3 月初发布。

首先&＃xff0c;PyTorch 2.0 引入了 torch.compile&＃xff0c;这是一种编译模式&＃xff0c;可以在不更改模型代码的情况下加速模型。在 163 个涵盖视觉、NLP 和其他领域的开源模型中&＃xff0c;该团队发现使用 2.0 可以将训练速度提高 38-76%。

其次&＃xff0c;PyTorch 2.0 是 100% 向后兼容的&＃xff1a;代码库一样&＃xff0c;API 一样&＃xff0c;写模型的方式也一样。团队之所以称它为 2.0&＃xff0c;是因为它有一些标志性的新特性&＃xff0c;包括&＃xff1a;

• TorchDynamo 可以从字节码分析生成 FX 图&＃xff1b;

• AOTAutograd 可以以 ahead-of-time 的方式生成反向图&＃xff1b;

• PrimTorch 引入了一个小型算子集&＃xff0c;使后端更容易&＃xff1b;

• TorchInductor&＃xff1a;一个由 OpenAI Triton 支持的 DL 编译器。

PyTorch 2.0 将延续 PyTorch 一贯的优势&＃xff0c;包括 Python 集成、命令式风格、API 简单等等。此外&＃xff0c;PyTorch 2.0 提供了相同的 eager-mode 开发和用户体验&＃xff0c;同时从根本上改变和增强了 PyTorch 在编译器级别的运行方式。该版本能够为「Dynamic Shapes」和分布式运行提供更快的性能和更好的支持。

在官方博客中&＃xff0c;PyTorch团队还公布了他们对于整个2.0系列的展望&＃xff1a;

以下是详细内容。

PyTorch 2.X&＃xff1a;速度更快、更加地 Python 化、一如既往地 dynamic

PyTorch 2.0 官宣了一个重要特性——torch.compile&＃xff0c;这一特性将 PyTorch 的性能推向了新的高度&＃xff0c;并将 PyTorch 的部分内容从 C&＃43;&＃43; 移回 Python。torch.compile 是一个完全附加的&＃xff08;可选的&＃xff09;特性&＃xff0c;因此 PyTorch 2.0 是 100% 向后兼容的。

支撑 torch.compile 的技术包括研发团队新推出的 TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。

• TorchDynamo 使用 Python Frame Evaluation Hooks 安全地捕获 PyTorch 程序&＃xff0c;这是一项重大创新&＃xff0c;是研究团队对快速可靠地获取图进行 5 年研发的结果&＃xff1b;

• AOTAutograd 重载 PyTorch 的 autograd 引擎作为一个跟踪 autodiff&＃xff0c;用于生成 ahead-of-time 向后跟踪&＃xff1b;

• PrimTorch 将约 2000 多个 PyTorch 算子规范化为一组约 250 个原始算子的闭集&＃xff0c;开发人员可以将其作为构建完整 PyTorch 后端的目标。这大大降低了编写 PyTorch 特性或后端的障碍&＃xff1b;

• TorchInductor 是一种深度学习编译器&＃xff0c;可为多个加速器和后端生成快速代码。对于 NVIDIA GPU&＃xff0c;它使用 OpenAI Triton 作为关键构建块。

TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 是用 Python 编写的&＃xff0c;并支持 dynamic shapes&＃xff08;即能够发送不同大小的张量而无需重新编译&＃xff09;&＃xff0c;这使得它们具备灵活、易于破解的特性&＃xff0c;降低了开发人员和供应商的使用门槛。

为了验证这些技术&＃xff0c;研发团队在各种机器学习领域测试了 163 个开源模型。实验精心构建了测试基准&＃xff0c;包括各种 CV 任务&＃xff08;图像分类、目标检测、图像生成等&＃xff09;、NLP 任务&＃xff08;语言建模、问答、序列分类、推荐系统等&＃xff09;和强化学习任务&＃xff0c;测试模型主要有 3 个来源&＃xff1a;

• 46 个来自 HuggingFace Transformers 的模型&＃xff1b;

• 来自 TIMM 的 61 个模型&＃xff1a;一系列 SOTA PyTorch 图像模型&＃xff1b;

• 来自 TorchBench 的 56 个模型&＃xff1a;包含来自 github 的精选流行代码库。

然后研究者测量加速性能并验证这些模型的准确性。加速可能取决于数据类型&＃xff0c;研究团队选择测量 float32 和自动混合精度 (AMP) 的加速。

在 163 个开源模型中&＃xff0c;torch.compile 在 93% 的情况下都有效&＃xff0c;模型在 NVIDIA A100 GPU 上的训练速度提高了 43%。在 float32 精度下&＃xff0c;它的平均运行速度提高了 21%&＃xff0c;而在 AMP 精度下&＃xff0c;它的运行速度平均提高了 51%。

目前&＃xff0c;torch.compile 还处于早期开发阶段&＃xff0c;预计 2023 年 3 月上旬将发布第一个稳定的 2.0 版本。

TorchDynamo&＃xff1a;快速可靠地获取图

TorchDynamo 是一种使用 Frame Evaluation API &＃xff08;PEP-0523 中引入的一种 CPython 特性&＃xff09;的新方法。研发团队采用数据驱动的方法来验证其在 Graph Capture 上的有效性&＃xff0c;并使用 7000 多个用 PyTorch 编写的 Github 项目作为验证集。TorchScript 等方法大约在 50% 的时间里都难以获取图&＃xff0c;而且通常开销很大&＃xff1b;而 TorchDynamo 在 99% 的时间里都能获取图&＃xff0c;方法正确、安全且开销可忽略不计&＃xff08;无需对原始代码进行任何更改&＃xff09;。这说明 TorchDynamo 突破了多年来模型权衡灵活性和速度的瓶颈。

TorchInductor&＃xff1a;使用 define-by-run IR 快速生成代码

对于 PyTorch 2.0 的新编译器后端&＃xff0c;研发团队从用户编写高性能自定义内核的方式中汲取灵感&＃xff1a;越来越多地使用 Triton 语言。此外&＃xff0c;研究者还想要一个编译器后端——使用与 PyTorch eager 类似的抽象&＃xff0c;并且具有足够的通用性以支持 PyTorch 中广泛的功能。

TorchInductor 使用 pythonic define-by-run loop level IR 自动将 PyTorch 模型映射到 GPU 上生成的 Triton 代码和 CPU 上的 C&＃43;&＃43;/OpenMP。TorchInductor 的 core loop level IR 仅包含约 50 个算子&＃xff0c;并且是用 Python 实现的&＃xff0c;易于破解和扩展。

AOTAutograd&＃xff1a;将 Autograd 重用于 ahead-of-time 图

PyTorch 2.0 的主要特性之一是加速训练&＃xff0c;因此 PyTorch 2.0 不仅要捕获用户级代码&＃xff0c;还要捕获反向传播。此外&＃xff0c;研发团队还想要复用现有的经过实践检验的 PyTorch autograd 系统。AOTAutograd 利用 PyTorch 的 torch_dispatch 可扩展机制来跟踪 Autograd 引擎&＃xff0c;使其能够「ahead-of-time」捕获反向传递&＃xff08;backwards pass&＃xff09;。这使 TorchInductor 能够加速前向和反向传递。

PrimTorch&＃xff1a;稳定的原始算子

为 PyTorch 编写后端具有挑战性。PyTorch 有 1200 多个算子&＃xff0c;如果考虑每个算子的各种重载&＃xff0c;则有 2000 多个。

在 PrimTorch 项目中&＃xff0c;研发团队致力于定义更小且稳定的算子集&＃xff0c;将 PyTorch 程序缩减到这样较小的算子集。目标是定义两个算子集&＃xff1a;

• Prim ops&＃xff1a;约有 250 个相当低级的算子。这些算子适用于编译器&＃xff0c;需要将它们重新融合在一起以获得良好的性能&＃xff1b;

• ATen ops&＃xff1a;约有 750 个规范算子。这些算子适用于已经在 ATen 级别集成的后端或没有编译功能的后端&＃xff08;无法从较低级别的算子集&＃xff08;如 Prim ops&＃xff09;恢复性能&＃xff09;。

用户体验

PyTorch 2.0 引入了一个简单的函数 torch.compile&＃xff0c;它会返回一个编译后的模型。

compiled_model *&＃61;* torch.compile(model)

compiled_model 保存对模型的引用&＃xff0c;并将 forward 函数编译为一个更优化的版本。在编译模型时&＃xff0c;PyTorch 2.0 给了几项设置来调整它&＃xff1a;

def torch.compile(model: Callable,*,mode: Optional[str] &＃61; "default",dynamic: bool &＃61; False,fullgraph:bool &＃61; False,backend: Union[str, Callable] &＃61; "inductor",# advanced backend options go here as kwargs**kwargs) -> torch._dynamo.NNOptimizedModule

• 「mode」指定编译器在编译时应该优化的内容。

• default 是一种预设模式&＃xff0c;它试图在不花费太长时间或使用额外内存的情况下高效编译。

• 其他模式&＃xff0c;如 reduce-overhead&＃xff0c;可以大大降低框架开销&＃xff0c;但要消耗少量额外内存。max-autotune 编译很长时间&＃xff0c;试图为你提供它所能生成的最快的代码。

• 「dynamic」模式指定是否为 Dynamic Shapes 启用代码路径。某些编译器优化不能应用于动态形状的程序。明确你想要一个带有动态形状还是静态形状的编译程序&＃xff0c;将有助于编译器提供更好的优化代码。

• 「fullgraph」类似于 Numba 的 nopython。它将整个程序编译成一个图&＃xff0c;或者给出一个错误提示&＃xff0c;解释为什么它不能这样做。大多数用户不需要使用这种模式。如果你非常注重性能&＃xff0c;那么你可以尝试使用它。

• 「backend 」指定使用哪个编译器后端。默认情况下使用 TorchInductor&＃xff0c;但还有其他一些可用的工具。

编译体验想要在默认模式中提供最大的好处和最大的灵活性。

常见问答

1、什么是 PT 2.0&＃xff1f;

2.0 是最新的 PyTorch 版本。PyTorch 2.0 提供相同的 eager 模式开发体验&＃xff0c;同时通过 torch.compile 添加编译模式。这种编译模式有可能在训练和推理期间加速模型。

2、为什么是 2.0 而不是 1.14&＃xff1f;

PyTorch 2.0 就是计划中的 1.14 。我们发布了大量新特性&＃xff0c;我们相信这些新特性会改变用户使用 PyTorch 的方式&＃xff0c;因此称其为 PyTorch 2.0。

3、如何安装 2.0&＃xff1f;有什么额外要求吗&＃xff1f;

安装最新的 nightlies&＃xff1a;

CUDA 11.7

pip3 install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117

CUDA 11.6

pip3 install numpy --pre torch[dynamo] torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu116

CPU

pip3 install numpy --pre torch torchvision torchaudio --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cpu

4、2.0 代码是否向后兼容 1.X&＃xff1f;

是的&＃xff0c;使用 2.0 不需要修改 PyTorch 工作流程。一行代码 model &＃61; torch.compile(model) 就可以优化模型&＃xff0c;以使用 2.0 堆栈&＃xff0c;并与其余 PyTorch 代码一起顺利运行。这是一个可选择项&＃xff0c;不需要使用新的编译器。

5、默认启用 2.0 吗&＃xff1f;

不&＃xff0c;必须通过使用单个函数调用优化模型从而在 PyTorch 代码中显式启用 2.0。

6、如何将 PT1.X 代码迁移到 PT2.0&＃xff1f;

代码应该按原样工作&＃xff0c;无需任何迁移。如果想使用 2.0 中引入的新编译模式特性&＃xff0c;那么可以从优化模型开始&＃xff1a;

model &＃61; torch.compile(model)

虽然加速主要是在训练期间观察到的&＃xff0c;但如果你的模型运行速度比 eager 模式快&＃xff0c;也可以将它用于推理。

import torchdef train(model, dataloader):model &＃61; torch.compile(model)for batch in dataloader:run_epoch(model, batch)def infer(model, input):model &＃61; torch.compile(model)return model(\*\*input)

7、是否有任何不该使用 PT 2.0 的应用程序&＃xff1f;

当前版本的 PT 2.0 仍处于实验阶段&＃xff0c;并且处于 nightlies 版本。其中的动态形状支持还处于早期阶段&＃xff0c;所以可以等到 2023 年 3 月稳定版发布后再使用该功能。

8、运行 PyTorch 2.0 时&＃xff0c;代码有何不同&＃xff1f; 

开箱即用&＃xff0c;PyTorch 2.0 与 PyTorch 1.x 相同&＃xff0c;模型以 eager 模式运行&＃xff0c;即 Python 的每一行都逐个执行。

在 2.0 中&＃xff0c;如果用 model &＃61; torch.compile(model) 将模型打包&＃xff0c;则模型在执行之前会经过 3 个步骤&＃xff1a;

1&＃xff09;图获取&＃xff1a;首先将模型重写为子图块。可由 TorchDynamo 编译的子图被「压平」&＃xff0c;其他子图&＃xff08;可能包含控制流代码或其他不受支持的 Python 结构&＃xff09;将回退到 Eager 模式。

2&＃xff09;Graph lowering&＃xff1a;所有 PyTorch 操作都被分解为特定于所选后端的组成内核。

3&＃xff09;图编译&＃xff0c;内核调用其相应的低级设备专用操作。

9、2.0 目前支持哪些编译器后端&＃xff1f;

默认和最完整的后端是 TorchInductor&＃xff0c;但是 TorchDynamo 有一个不断增长的后端列表&＃xff0c;可以通过调用 torchdynamo.list_backends(). 找到

10、2.0 版本的分布式训练能力如何&＃xff1f;

Compiled 模式下的 DDP 和 FSDP &＃xff0c;比 FP32 中的 Eager 模式快 15%、AMP 精度快 80%。PT2.0 做了一些额外的优化&＃xff0c;以确保 DDP 的通信 - 计算 overlap 与 Dynamo 的部分图创建良好协作。想要确保使用 static_graph&＃61;False 运行 DDP&＃xff0c;更多细节参见&＃xff1a;https://dev-discuss.pytorch.org/t/torchdynamo-update-9-making-ddp-work-with-torchdynamo/860

11、为什么我的代码用 2.0 的 Compiled Model 运行变慢&＃xff1f;

性能下降最可能的原因是 graph break 太多。例如&＃xff0c;类似模型前向 trigger 中的输出语句这样的东西会触发 graph break。详见&＃xff1a;https://pytorch.org/docs/master/dynamo/faq.html#why-am-i-not-seeing-speedups

12、以前运行的代码在 2.0 中崩溃了&＃xff0c;该如何调试&＃xff1f;

参见&＃xff1a;https ://pytorch.org/docs/master/dynamo/faq.html#why-is-my-code-crashing

参考链接&＃xff1a;https://pytorch.org/get-started/pytorch-2.0/

点击进入—> CV 微信技术交流群

CVPR/ECCV 2022论文和代码下载

后台回复&＃xff1a;CVPR2022&＃xff0c;即可下载CVPR 2022论文和代码开源的论文合集

后台回复&＃xff1a;ECCV2022&＃xff0c;即可下载ECCV 2022论文和代码开源的论文合集

后台回复&＃xff1a;Transformer综述&＃xff0c;即可下载最新的3篇Transformer综述PDF

目标检测和Transformer交流群成立
扫描下方二维码&＃xff0c;或者添加微信&＃xff1a;CVer222&＃xff0c;即可添加CVer小助手微信&＃xff0c;便可申请加入CVer-目标检测或者Transformer 微信交流群。另外其他垂直方向已涵盖&＃xff1a;目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer等。
一定要备注&＃xff1a;研究方向&＃43;地点&＃43;学校/公司&＃43;昵称&＃xff08;如目标检测或者Transformer&＃43;上海&＃43;上交&＃43;卡卡&＃xff09;&＃xff0c;根据格式备注&＃xff0c;可更快被通过且邀请进群▲扫码或加微信号: CVer222&＃xff0c;进交流群
CVer学术交流群&＃xff08;知识星球&＃xff09;来了&＃xff01;想要了解最新最快最好的CV/DL/ML论文速递、优质开源项目、学习教程和实战训练等资料&＃xff0c;欢迎扫描下方二维码&＃xff0c;加入CVer学术交流群&＃xff0c;已汇集数千人&＃xff01;▲扫码进群
▲点击上方卡片&＃xff0c;关注CVer公众号

整理不易&＃xff0c;请点赞和在看