内容一览： 昨夜召开的 PyTorch Conference 2022 中，官方正式发布了 PyTorch 2.0。本文将收拾 PyTorch 2.0 与 1.x 比较的最大差异。

关键词： PyTorch 2.0 编译器 机器学习

本文首发自微信大众号：HypeAI超神经

在 PyTorch Conference 2022 上，PyTorch 官方正式发布了 PyTorch 2.0，整场活动含「compiler」率极高，跟从前的 1.x 版别比较，2.0 中有了「推翻式」的改变。

PyTorch 2.0 中发布了很多足以改变 PyTorch 运用办法的新功用，它供给了相同的 eager mode 和用户体验，一起通过 torch.compile 增加了一个编译形式， 在练习和推理过程中能够对模型进行加快，从而供给更佳的功用和对 Dynamic Shapes 及 Distributed 的支撑。

本文将对 PyTorch 2.0 进行具体简介。

太长不看版

  • PyTorch 2.0 在保留原有优势的一起，大举支撑编译
  • torch.compile 为可选功用，只需一行代码即可运转编译
  • 4 项重要技能：TorchDynamo、AOTAutograd、PrimTorch 以及 TorchInductor
  • 5 年前就尝试过编译，效果一向不抱负
  • PyTorch 1.x 代码无需向 2.0 搬迁* PyTorch 2.0 stable 版别预计明年 3 月发布

更快、更优、编译支撑

在昨夜的 PyTorch Conference 2022 中，官方正式发布了 torch.compile， 它使得 PyTorch 的功用进一步提高，并开始将 PyTorch 的部分内容从 C++ 中回到 Python。

PyTorch 2.0 中的最新技能包括：

TorchDynamo、AOTAutograd、PrimTorch 以及 TorchInductor。

1. TorchDynamo

它能够凭借 Python Frame Evaluation Hooks，安全地获取 PyTorch 程序，这项严重立异是 PyTorch 曩昔 5 年来在安全图结构捕获 (safe graph capture) 方面的研发效果汇总。

2. AOTAutograd

重载 PyTorch autograd engine，作为一个 tracing autodiff，用于生成超前的 backward trace。

3. PrimTorch

将 2000+ PyTorch 算子归纳为约 250 个 primitive operator 闭集 (closed set)，开发者能够针对这些算子构建一个完好的 PyTorch 后端。PrimTorch 大大简化了编写 PyTorch 功用或后端的流程。

4. TorchInductor一个深度学习编译器，能够为多个加快器和后端生成 fast code。关于 NVIDIA GPU，它运用 OpenAI Triton 作为关键构建模块。

TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 是用 Python 编写的， 并且支撑 dynamic shape（无需从头编译就能发送不同巨细的向量），这使得它们灵敏且易学，降低了开发者和供应商的准入门槛。

为了验证这些技能，PyTorch 官方运用了机器学习范畴的 163 个开源模型， 包括图画分类、方针检测、图画生成等使命，以及各种 NLP 使命，如言语建模、问答、序列分类、引荐系统和强化学习。这些 Benchmark 分为三类：

  • 来自 HuggingFace Transformers 的 46 个模型
  • 来自 TIMM 的 61 个模型：由 Ross Wightman 搜集的 SoTA PyTorch 图画模型
  • 来自 TorchBench 的 56 个模型：GitHub 上搜集的一组盛行代码库。

关于开源模型，PyTorch 官方没有进行修正，仅仅增加了一个 torch.compile 调用来进行封装。

接下来 PyTorch 工程师在这些模型中丈量速度并验证精度，由于提速或许取决于数据类型，因而官方在 float32 和主动混合精度 (AMP) 上都丈量了提速。 由于 AMP 在实践中更常见，测试比例设定为：0.75 * AMP + 0.25 * float32 的。

在这 163 个开源模型中，torch.compile 能够在 93% 模型上正常运转， 运转往后，模型在 NVIDIA A100 GPU 上的运转速度达到了 43% 的提高。在 Float32 精度下，运转速度均匀提高 21%；在 AMP 精度下，运转速度均匀提高 51%。

注意：在桌面级 GPU（如 NVIDIA 3090）上，测得的速度比在服务器级 GPU（如 A100）上要低。截至现在，PyTorch 2.0 默许后端 TorchInductor 现已支撑 CPU 和 NVIDIA Volta 和 Ampere GPU，暂不支撑其他 GPU、xPU 或更老的 NVIDIA GPU。

PyTorch 2.0 重磅发布：编译、编译、还是编译！
NVIDIA A100 GPU eager mode torch.compile 针对不同模型的提速体现

torch.compile 在线试用： 开发者能够通过 nightly 二进制文件安装试用，PyTorch 2.0 Stable 版别预计将于 2023 年 3 月初发布。

在 PyTorch 2.x roadmap 中，编译形式 (compiled mode) 的功用和可扩展性在未来会不断进行丰厚和提高。

PyTorch 2.0 重磅发布：编译、编译、还是编译！
PyTorch 2.x roadmap

​

开发布景

PyTorch 的开发理念自始至终都是 flexibility 和 hackability 第一，功用则是第二， 致力于：

1. 高功用的 eager execution

2. 不断 Python 化内部结构

3. Distributed, Autodiff, Data loading, Accelerators 等的杰出抽象

PyTorch自2017年问世以来，硬件加快器（如GPU）的计算速度提高了约 15倍，内存访问速度提高了约 2 倍。

为了保持高功用的 eager execution，PyTorch 内部的大部分内容不得不转移到 C++ 中，这使得 PyTorch hackability 下降，也增加了开发者参与代码贡献的门槛。

从第一天起，PyTorch 官方就认识到了 eager execution 的功用局限。2017 年 7 月，官方开始致力于为 PyTorch 开发一个编译器。该编译器需求在不献身 PyTorch 体验的前提下，加快 PyTorch 程序的运转，其关键标准是保持某种程度上的灵敏性 (flexibility)：支撑开发者广泛运用的 dynamic shapes 以及 dynamic programs。

PyTorch 2.0 重磅发布：编译、编译、还是编译！

PyTorch 技能细节

自问世以来，PyTorch 中建立过好几个编译器项目，这些编译器能够分为 3 类：

  • 图结构的获取 (graph acquisition)
  • 图结构的降低 (graph lowering)
  • 图结构的编译 (graph compilation)

其中，图结构的获取面对的挑战最多。

曩昔5年中，官方尝试了 torch.jit.trace、TorchScript、FX tracing 以及 Lazy Tensors，但它们有些够灵敏但不够快，有些够快但不灵敏，有些既不快也不灵敏，有些用户体验不好。

尽管 TorchScript 很有前途，但它需求很多修正代码和依靠，可行性并不高。

PyTorch 2.0 重磅发布：编译、编译、还是编译！
PyTorch 编译流程示意图

​

TorchDynamo：牢靠快速地获取图结构

TorchDynamo 运用了 PEP-0523 中引进的CPython 功用，称为结构评估 API (Frame Evaluation API)。官方采取了一种数据驱动的办法来验证其在 Graph Capture 上的有效性，运用 7000 多个用 PyTorch 编写的 Github 项目作为验证集。

实验表明，TorchDynamo 在 99% 的时间里都能正确、安全地获取图结构，并且开销能够忽略不计， 由于它无需对原始代码做任何修正。

TorchInductor：用 define-by-run IR 进行更敏捷的 codegen

越来越多的开发者在编写高功用自界说内核时，会运用 Triton 言语。 此外，关于 PyTorch 2.0 全新的编译器后端，官方还期望能够运用与 PyTorch eager 类似的抽象，并且具有满足的通用功用支撑 PyTorch 中广泛的功用。

TorchInductor 运用 Pythonic define-by-run loop level IR，主动将 PyTorch 模型映射到 GPU 上生成的 Triton 代码以及 CPU 上的 C++/OpenMP。

TorchInductor 的 core loop level IR 只包括大约 50 个算子，并且是用 Python 实现的，这使得它具有很强的 hackability 和扩展性。

AOTAutograd：关于 ahead-of-time graph，重用 Autograd

PyTorch 2.0 要想加快练习，不仅要捕获用户级代码，并且要捕获反向传达算法 (backpropagation)。假如能用上通过验证的 PyTorch autograd system 就更好了。

AOTAutograd 利用 PyTorch torch_dispatch 扩展机制来追寻 Autograd engine， 使开发者得以「ahead-of-time」捕获反向传达 (backwards pas)，从而使开发者得以运用 TorchInductor 加快 forwards 和 backwards pass。

PrimTorch：安稳的 Primitive operator

为 PyTorch 写一个后端并不简单，Torch 有 1200+ 算子，假如考虑到每个算子的各种重载 (overload)，数量高达 2000+。

PyTorch 2.0 重磅发布：编译、编译、还是编译！

2000+ PyTorch 算子的分类概况

因而，编写后端或交叉功用 (cross-cutting feature) 成为一项消耗精力的作业。PrimTorch 致力于界说更小更安稳的算子集。 PyTorch 程序能够继续降级 (lower) 到这些算子集。官方的方针是界说两个算子集：

  • Prim ops 包括约 250 个相对底层的算子，由于满足底层，所以这些算子更适用于编译器，开发者需求将这些算子进行交融，才干获得杰出的功用。
  • ATen ops 包括约 750 个典型算子 (canonical operator)，适合于直接输出。这些算子适用于现已在 ATen 等级上集成的后端，或许没有通过编译的后端，才干从底层算子集（如 Prim ops) 恢复功用。

FAQ

1. 怎么安装 PyTorch 2.0？有哪些额定要求？

安装最新的 nightlies：

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

2. PyTorch 2.0 的代码是否向下兼容 1.x？

是的，2.0 不要求修正 PyTorch workflow，只需一行代码

model = torch.compile(model)

即可优化模型运用 2.0 stack，并与 PyTorch 其他代码顺畅运转。该选项不强制，开发者仍可运用从前的版别。

3. PyTorch 2.0 是否默许启用？

不是，必须在 PyTorch 代码中清晰启用 2.0，办法是通过一个单一函数调用 (single function call) 来优化模型。

4. 怎么将 PT1.X 代码搬迁到 PT2.0？

从前的代码不需求任何搬迁，假如想运用 2.0 中引进的全新的 compiled mode 功用，能够先用一行代码来优化模型：

model = torch.compile(model)

速度提高主要体现在练习过程中，假如模型运转速度快于 eager mode，则表示能够用于推理。

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)

5. PyTorch 2.0 中哪些功用不引荐运用 ？

现在 PyTorch 2.0 没有安稳，且仍处于 nightlies 版别。torch.compile 中的对 Dynamic shapes 的支撑仍处于前期，在 2023 年 3 月安稳版 2.0 发布前不主张运用。

也就是说，即使是 static-shaped workload，仍然在构建 compiled mode，或许会呈现一些 bug。关于代码中呈现崩溃的部分，请禁用 compiled mode 并提交一个 issue。

提交 issue 传送门：

github.com/pytorch/pyt…

以上就是关于PyTorch2.0的具体介绍，后续我们将完好收拾 PyTorch 2.0 Get Started 介绍，欢迎继续重视我们！

您也能够微信查找 Hyperai01 ，和神经星星一起进入 PyTorch 技能开发群评论。