作者：PyTorch 团队

我们很高兴宣布 PyTorch 1.8 版本现已可用。此版本包含自 1.7 版本以来的 3000 多次提交。它包括编译器、代码优化、科学计算的前端 API 以及通过 pytorch.org 提供的二进制文件支持 AMD ROCm 的重大更新和新功能。它还为管道并行和模型并行以及梯度压缩提供了大型训练的改进功能。以下是一些亮点：

  1. 通过 torch.fx 支持 Python 到 Python 的函数式转换；
  2. 添加或稳定了 API 以支持 FFT（torch.fft）、线性代数函数（torch.linalg），添加了对复杂张量自动求导的支持，并更新以提高计算 Hessian 和 Jacobian 的性能；以及
  3. 分布式训练的重大更新和改进，包括：提高了 NCCL 的可靠性；支持管道并行；RPC 性能分析；以及通过在 DDP 中添加通信钩子支持梯度压缩。请在此处查看完整的发布说明。

伴随 1.8 版本，我们还发布了 PyTorch 库的重大更新，包括 TorchCSPRNG、TorchVision、TorchText 和 TorchAudio。有关库发布的更多信息，请在此处查看博文。如前所述，PyTorch 发布中的功能分为稳定版 (Stable)、测试版 (Beta) 和原型版 (Prototype)。您可以在此处的博文中了解更多定义。

新增和更新的 API

PyTorch 1.8 版本带来了一系列新增和更新的 API 界面，从与 NumPy 兼容的额外 API，到提高代码在推理和训练时的性能并对其进行扩展的方式。以下是此版本中主要功能的简要总结：

[稳定版] Torch.fft 支持高性能 NumPy 风格的 FFT

作为 PyTorch 支持科学计算目标的一部分，我们投入了精力改进我们的 FFT 支持，并在 PyTorch 1.8 版本中发布了 torch.fft 模块。此模块实现了与 NumPy 的 np.fft 模块相同的功能，但支持硬件加速和自动求导。

[测试版] 通过 torch.linalg 支持 NumPy 风格的线性代数函数

torch.linalg 模块以 NumPy 的 np.linalg 模块为蓝本，为常见的线性代数操作（包括 Cholesky 分解、行列式、特征值等等）带来了 NumPy 风格的支持。

[测试版] 使用 FX 进行 Python 代码转换

FX 允许您编写 transform(input_module : nn.Module) -> nn.Module 形式的转换，您可以将 Module 实例输入，并从中获得转换后的 Module 实例。

这类功能适用于许多场景。例如，基于 FX 的图模式量化产品与 FX 同时作为原型发布。图模式量化通过利用 FX 的程序捕获、分析和转换功能，自动化了神经网络量化过程。我们还在开发许多其他基于 FX 的转换产品，我们很高兴与社区分享这个强大的工具包。

由于 FX 转换使用并生成 nn.Module 实例，它们可以在许多现有 PyTorch 工作流程中使用。这包括例如在 Python 中训练然后通过 TorchScript 部署的工作流程。

您可以在官方文档中阅读更多关于 FX 的信息。您还可以在此处找到几个使用 torch.fx 实现的程序转换示例。我们一直在改进 FX，并邀请您在论坛或问题跟踪器上分享您对该工具包的任何反馈。

我们要感谢 TorchScript 追踪、Apache MXNet hybridize，以及最近的 JAX 作为通过追踪进行程序获取的影响因素。我们还要感谢 Caffe2、JAX 和 TensorFlow 在简单、定向数据流图程序表示及其转换价值方面的启发。

分布式训练

PyTorch 1.8 版本添加了许多新功能，并改进了可靠性和可用性。具体来说，支持：增加了稳定级别的异步错误/超时处理以提高 NCCL 可靠性；以及基于 RPC 的性能分析的稳定支持。此外，我们还添加了对管道并行性的支持，以及通过在 DDP 中使用通信钩子实现的梯度压缩。详情如下：

[测试版] 管道并行

随着机器学习模型尺寸不断增长，传统的分布式数据并行 (DDP) 训练不再适用，因为这些模型无法完全放入单个 GPU 设备。新的管道并行功能提供了易于使用的 PyTorch API，以便在训练循环中利用管道并行。

[测试版] DDP 通信钩子

DDP 通信钩子是一个通用接口，通过覆盖 DistributedDataParallel 中的普通 allreduce 操作来控制如何在 worker 之间通信梯度。提供了一些内置的通信钩子，包括 PowerSGD，用户可以轻松应用其中任何一个钩子来优化通信。此外，通信钩子接口还可以支持用户自定义的通信策略以应对更高级的用例。

分布式训练的其他原型功能

除了此版本中主要的稳定版和测试版分布式训练功能外，我们还在夜间构建版本中提供了一些原型功能供大家试用并提供反馈。以下链接了文档草案供参考：

  • (原型) ZeroRedundancyOptimizer - 基于并与 Microsoft DeepSpeed 团队合作开发，此功能通过在 ProcessGroup 组中的所有参与进程之间分片优化器状态来帮助减少每个进程的内存占用。有关更多详细信息，请参阅此文档。
  • (原型) Process Group NCCL Send/Recv - NCCL send/recv API 在 v2.7 中引入，此功能在 NCCL 进程组中添加了对其的支持。此功能将为用户提供一个选项，以便在 Python 层而不是 C++ 层实现集合操作。要了解更多信息，请参阅此文档和代码示例。
  • (原型) 使用 TensorPipe 在 RPC 中支持 CUDA - 此功能将为使用 PyTorch RPC 的多 GPU 机器用户带来显著的速度提升，因为 TensorPipe 在可用时将自动利用 NVLink，并在进程之间交换 GPU 张量时避免昂贵的 host 内存复制。如果不在同一台机器上，TensorPipe 将回退到将张量复制到 host 内存并将其作为普通 CPU 张量发送。这也将改善用户体验，因为用户将能够在代码中像处理普通 CPU 张量一样处理 GPU 张量。有关更多详细信息，请参阅此文档。
  • (原型) Remote Module - 此功能允许用户像使用本地模块一样在远程 worker 上操作模块，其中 RPC 对用户是透明的。过去，此功能是临时实现的，总的来说，此功能将提高 PyTorch 上模型并行性的可用性。有关更多详细信息，请参阅此文档。

PyTorch 移动端

PyTorch 移动端的支持正在扩展，提供了一组新的教程，帮助新用户更快地在设备上部署模型，并为现有用户提供工具，以便从我们的框架中获得更多收益。这包括：

我们新的演示应用还包括图像分割、目标检测、神经机器翻译、问答和视觉 transformer 等示例。它们在 iOS 和 Android 上均可用。

除了 MobileNetV3 和其他模型在 CPU 上的性能改进外，我们还改进了 Android GPU 后端原型，以支持更广泛的模型并加快推理速度。

最后，在此版本中，我们推出了 PyTorch Mobile Lite Interpreter 作为原型功能。Lite Interpreter 允许用户减小运行时二进制文件大小。请尝试使用并向我们提供您在 PyTorch 论坛上的反馈。我们的所有最新更新都可以在 PyTorch 移动端页面找到。

[原型] PyTorch Mobile Lite Interpreter

PyTorch Lite Interpreter 是 PyTorch 运行时的一个精简版本，可以在资源受限的设备上执行 PyTorch 程序，并减小二进制文件大小。与当前版本中的设备上运行时相比，此原型功能可将二进制文件大小减小高达 70%。

性能优化

在 1.8 版本中，我们发布了对基准测试工具的支持，使用户能够更好地监控性能。我们还开放了一个新的自动化量化 API。详情如下：

(测试版) 基准测试工具

基准测试工具允许用户进行准确的性能测量，并提供可组合的工具，帮助进行基准测试制定和后处理。这对于 PyTorch 贡献者来说应该很有帮助，可以快速了解他们的贡献对 PyTorch 性能的影响。

示例

from torch.utils.benchmark import Timer

results = []
for num_threads in [1, 2, 4]:
    timer = Timer(
        stmt="torch.add(x, y, out=out)",
        setup="""
            n = 1024
            x = torch.ones((n, n))
            y = torch.ones((n, 1))
            out = torch.empty((n, n))
        """,
        num_threads=num_threads,
    )
    results.append(timer.blocked_autorange(min_run_time=5))
    print(
        f"{num_threads} thread{'s' if num_threads > 1 else ' ':<4}"
        f"{results[-1].median * 1e6:>4.0f} us   " +
        (f"({results[0].median / results[-1].median:.1f}x)" if num_threads > 1 else '')
    )

1 thread     376 us   
2 threads    189 us   (2.0x)
4 threads     99 us   (3.8x)

(原型) FX 图模式量化

FX 图模式量化是 PyTorch 中新的自动化量化 API。它通过添加对函数式支持并自动化量化过程来改进了 Eager 模式量化，尽管用户可能需要重构模型以使其与 FX 图模式量化兼容（可以使用 torch.fx 进行符号追踪）。

硬件支持

[测试版] 在 C++ 中为新的后端扩展 PyTorch Dispatcher 的能力

在 PyTorch 1.8 中，您现在可以创建位于 pytorch/pytorch 仓库之外的新的 Out-of-tree 设备。以下链接的教程展示了如何注册您的设备并使其与原生的 PyTorch 设备保持同步。

[测试版] AMD GPU 二进制文件现已可用

从 PyTorch 1.8 版本开始，我们增加了对 ROCm wheel 的支持，为使用 AMD GPU 提供了简便的入门方式。您只需访问标准的 PyTorch 安装选择器，选择 ROCm 作为安装选项并执行提供的命令即可。

感谢您的阅读。如果您对这些更新感到兴奋并希望参与 PyTorch 的未来发展，我们鼓励您加入讨论论坛并在 GitHub 上提交 issues。

致敬！

PyTorch 团队