PyTorch 1.8发布，支持 AMD GPU 和 Python 函数转换

近日，PyTorch 团队发布了 PyTorch 1.8 版本。该版本整合了自去年 10 月 1.7 版本发布以来的 3000 多次 commit，提供了编译、代码优化、科学计算前端 API 方面的主要更新和新特性。值得一提的是，该版本新增了对 AMD ROCm 的支持。

此外，PyTorch 1.8 版本还为大规模训练 pipeline 和模型并行化、梯度压缩提供了特性改进。该版本的主要亮点如下：

• 支持 Python 函数转换；

• 添加或稳定化 API，以支持 FFT (torch.fft)、线性代数函数 (torch.linalg)；

• 添加对复杂张量 autograd 的支持；

• 多项更新用于提升 Hessian 与 Jacobian 矩阵计算的性能；

• 改进分布式训练，包括提升 NCCL 可靠性、支持 pipeline 并行化、RPC profiling，以及通过添加梯度压缩来支持通信钩子（hook）。

（注：从 PyTorch 1.6 版本开始，PyTorch 特性分为 Stable（稳定版）、Beta（测试版）和 Prototype（原型版）。

此外，PyTorch 团队还对多个 PyTorch 库进行了主要更新，包括 TorchCSPRNG、TorchVision、TorchText 和 TorchAudio。

新增和更新 API

通过 torch.fft 支持 NumPy 兼容的 FFT 操作

PyTorch 1.7 版本提出了这一特性的 Beta 版，而在 1.8 版本中该特性更新为稳定版。FFT 支持旨在完成 PyTorch 支持科学计算的目的。torch.fft 模块和 NumPy 的 np.fft 模块实现了同样的功能，并且支持硬件加速和 autograd。

通过 torch.linalg 支持 NumPy 式的线性代数函数

torch.linalg 模块类似于 NumPy 中的 np.linalg 模块，支持 NumPy 式的线性代数操作，包括 Cholesky 分解、行列式、特征值等。

使用 torch.fx 进行 Python 代码转换

这一 Beta 特性支持 Python 代码转换，开发者可以利用它做 Conv/BN 融合、图模式量化、实现 vmap 等。鉴于 torch.fx 提供 PyTorch 代码的图表示，开发者可以用 Python 写任意变换或分析。

分布式训练

pipeline 并行化

这一新增的 Beta 特性提供了一个易用的 PyTorch API，可将 pipeline 并行化作为训练 loop 的一部分。

DDP 通信钩子

DDP 通信钩子是一个通用接口，用于控制 workers 间的梯度通信。

此外，PyTorch 1.8 还增加了一些 prototype 特性：

• ZeroRedundancyOptimizer：有助于减少每个线程的内存占用；

• Process Group NCCL Send/Recv：该特性允许用户在 Python 层（而非 C++ 层）实现集合操作；

• CUDA-support in RPC using TensorPipe：该特性为使用 PyTorch RPC 和多 GPU 机器的用户带来速度提升；

• Remote Module：该特性允许用户像操作本地模块那样操作远程 worker 上的模块。

支持 PyTorch Mobile

此次版本更新发布了一组新的移动端教程，包括在 iOS 端和安卓端实现图像分割 DeepLabV3 模型。PyTorch 还发布了新的 demo app，包括图像分割、目标检测、神经机器翻译、问答和视觉 transformer。

此外，这次发布还包括 PyTorch Mobile Lite Interpreter，该解释器可降低运行时二进制文件大小。

性能优化

为了帮助用户更好地监控性能变化，PyTorch 1.8 版本支持 benchmark utils，并开放了新的自动量化 API——FX Graph Mode Quantization。

硬件支持

在硬件支持方面，PyTorch 1.8 版本新增了两个 Beta 特性：

扩展 PyTorch Dispatcher，使之适应新型 C++ 后端；

支持 AMD ROCm。

需要注意的是，PyTorch 1.8 仅在 Linux 系统中支持 AMD ROCm。

参考链接：

https://pytorch.org/blog/pytorch-1.8-released/

https://github.com/pytorch/pytorch

https://pytorch.org/

https://twitter.com/cHHillee/status/1367621538791317504