由 PyTorch 团队

我们很高兴地宣布 PyTorch 1.8 的发布。这个版本自 1.7 以来包含了 3000 多个提交。它包括编译、代码优化、科学计算的前端 API 以及通过 pytorch.org 提供的二进制文件对 AMD ROCm 的支持等主要更新和新功能。它还提供了针对大规模训练的改进功能，包括管道和模型并行性以及梯度压缩。其中一些亮点包括：

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

随着 1.8 版本的发布，我们还发布了 PyTorch 库的重大更新，包括 TorchCSPRNG、TorchVision、TorchText 和 TorchAudio。有关库发布的更多信息，请参阅此处帖子。如前所述，PyTorch 版本中的功能被分类为稳定版、测试版和原型版。您可以在此处帖子中了解更多关于定义的信息。

新增和更新 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 模块，为常见的线性代数操作提供了 NumPy 风格的支撑，包括 Cholesky 分解、行列式、特征值等。

[Beta] Python 代码变换与 FX

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，可以在训练循环中利用管道并行。

[Beta] DDP 通信钩子

DDP 通信钩子是一个通用的接口，用于通过覆盖 DistributedDataParallel 中的 vanilla allreduce 来控制如何在工作者之间通信梯度。提供了一些内置的通信钩子，包括 PowerSGD，用户可以轻松地将这些钩子应用于优化通信。此外，通信钩子接口还可以支持用户定义的通信策略，以适应更高级的使用场景。

分布式训练的附加原型功能

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

  • （原型）ZeroRedundancyOptimizer - 基于并与微软 DeepSpeed 团队合作，此功能通过在 ProcessGroup 团队中跨所有参与进程分片优化器状态，有助于减少每个进程的内存占用。请参阅以下文档以获取更多详细信息。
  • (原型) NCCL 进程组发送/接收 - NCCL 发送/接收 API 自 v2.7 版本引入，此功能为 NCCL 进程组添加了对它的支持。此功能将为用户提供在 Python 层而不是 C++层实现集体操作的选择。请参阅此文档和代码示例以获取更多信息。
  • (原型) 使用 TensorPipe 的 RPC 中的 CUDA 支持 - 此功能应会为使用 PyTorch RPC 的多 GPU 机器的用户带来显著的性能提升，因为 TensorPipe 将自动利用 NVLink（如果可用），并在交换进程间的 GPU 张量时避免昂贵的内存复制操作。当不在同一台机器上时，TensorPipe 将回退到将张量复制到主机内存，并以常规 CPU 张量的形式发送。这也会改善用户体验，因为用户可以在他们的代码中将 GPU 张量视为常规 CPU 张量。请参阅此文档以获取更多详细信息。
  • （原型）远程模块 - 该功能允许用户像使用本地模块一样操作远程工作器上的模块，其中 RPC 对用户是透明的。在过去，这种功能以临时方式实现，总体上这将提高 PyTorch 模型并行的可用性。有关详细信息，请参阅此文档。

PyTorch 移动

支持 PyTorch Mobile 正在扩展，新增一系列教程，帮助新用户更快地将模型部署到设备上，并为现有用户提供一个工具，以充分利用我们的框架。这些包括：

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

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

最后，我们在本次版本中作为原型功能推出了 PyTorch Mobile Lite 解释器。轻量级解释器允许用户减小运行时二进制文件大小。请尝试这些功能，并在 PyTorch 论坛上向我们发送您的反馈。所有最新更新都可以在 PyTorch Mobile 页面上找到。

[原型] PyTorch Mobile Lite 解释器

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

性能优化

在 1.8 版本中，我们将推出对基准工具的支持，以使用户能够更好地监控性能。我们还将开放一个新的自动化量化 API。以下是详细信息：

（Beta）基准工具

基准工具允许用户进行精确的性能测量，并提供可组合的工具来帮助进行基准测试的制定和后处理。这预计将有助于 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。它通过添加对函数的支持并自动化量化过程来改进即时模式量化，尽管人们可能需要重构模型以使模型与 FX 图模式量化（可使用 torch.fx 符号追踪）兼容。

硬件支持

[Beta] 能够在 C++ 中扩展 PyTorch 分派器以支持新的后端

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

[Beta] AMD GPU 可用二进制文件

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

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

喝彩！

PyTorch 团队