我们很高兴地宣布 PyTorch 1.8 已发布。自 1.7 版本以来,此版本包含了 3,000 多次提交。它包括编译、代码优化、科学计算前端 API 和通过 pytorch.org 提供的二进制文件支持 AMD ROCm 的重大更新和新功能。它还为管道和模型并行性以及梯度压缩提供了改进的大规模训练功能。以下是一些亮点:
- 通过
torch.fx
支持 Python 到 Python 的函数式转换; - 添加或稳定了用于支持 FFT (
torch.fft
)、线性代数函数 (torch.linalg
) 的 API,增加了对复杂张量自动微分的支持,并更新了以提高 Hessian 和 Jacobian 计算的性能;以及 - 分布式训练的重大更新和改进包括:改进了 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 模块,为常见的线性代数操作(包括 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 混合以及最近的 JAX 作为通过追踪进行程序获取的影响。我们还要感谢 Caffe2、JAX 和 TensorFlow 作为简单、有向数据流图程序表示及其转换的价值的灵感来源。
分布式训练
PyTorch 1.8 版本增加了一些新功能,并改进了可靠性和可用性。具体来说,增加了对:稳定级别的异步错误/超时处理以提高 NCCL 的可靠性;以及对 基于 RPC 的性能分析的稳定支持。此外,我们还增加了对管道并行性和通过 DDP 中的通信钩子进行梯度压缩的支持。详细信息如下:
[测试版] 管道并行性
随着机器学习模型规模的不断增长,传统的分布式数据并行 (DDP) 训练不再适用,因为这些模型无法在单个 GPU 设备上运行。新的管道并行功能提供了一个易于使用的 PyTorch API,可将管道并行作为训练循环的一部分。
[测试版] DDP 通信钩子
DDP 通信钩子是一个通用接口,通过覆盖 DistributedDataParallel 中的普通 allreduce 来控制如何跨 workers 传输梯度。提供了一些内置的通信钩子,包括 PowerSGD,用户可以轻松应用任何这些钩子来优化通信。此外,通信钩子接口还可以支持用户定义的通信策略,以应对更高级的用例。
分布式训练的其他原型功能
除了此版本中主要的稳定版和测试版分布式训练功能外,我们还有一些可在夜间版本中试用并提供反馈的原型功能。我们已在下面链接了草案文档供参考:
- (原型) ZeroRedundancyOptimizer – 基于 Microsoft DeepSpeed 团队并与其合作,此功能通过在
ProcessGroup
组中的所有参与进程之间分片优化器状态来帮助减少每个进程的内存占用。有关更多详细信息,请参阅此文档。 - (原型) 进程组 NCCL 发送/接收 – NCCL 发送/接收 API 在 v2.7 中引入,此功能为其在 NCCL 进程组中添加了支持。此功能将为用户提供一个选项,以在 Python 层而不是 C++ 层实现集体操作。有关更多信息,请参阅此文档和代码示例。
- (原型) RPC 中使用 TensorPipe 的 CUDA 支持 – 此功能应为使用 PyTorch RPC 的多 GPU 机器用户带来显著的速度提升,因为 TensorPipe 将在可用时自动利用 NVLink,并避免在进程之间交换 GPU 张量时 costly 的主机内存复制。当不在同一台机器上时,TensorPipe 将回退到将张量复制到主机内存并将其作为常规 CPU 张量发送。这也将改善用户体验,因为用户将能够在代码中将 GPU 张量视为常规 CPU 张量。有关更多详细信息,请参阅此文档。
- (原型) 远程模块 – 此功能允许用户在远程 worker 上操作模块,就像使用本地模块一样,其中 RPC 对用户是透明的。过去,此功能是以临时方式实现的,总体而言,此功能将改善 PyTorch 上模型并行性的可用性。有关更多详细信息,请参阅此文档。
PyTorch Mobile
PyTorch Mobile 的支持正在扩展,并推出了一系列新教程,以帮助新用户更快地在设备上启动模型,并为现有用户提供一个工具,以充分利用我们的框架。这些包括:
我们的新演示应用程序还包括图像分割、目标检测、神经机器翻译、问答和视觉转换器的示例。它们在 iOS 和 Android 上均可用。
除了 MobileNetV3 和其他模型在 CPU 上的性能改进外,我们还改进了 Android GPU 后端原型,以实现更广泛的模型覆盖和更快的推理。
最后,我们在此版本中推出了 PyTorch Mobile Lite Interpreter 作为原型功能。Lite Interpreter 允许用户减小运行时二进制文件的大小。请试用这些功能,并在PyTorch 论坛上向我们发送您的反馈。我们所有的最新更新都可以在PyTorch Mobile 页面上找到。
[原型] 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
存储库之外的新树外设备。下面链接的教程演示了如何注册您的设备并使其与原生 PyTorch 设备保持同步。
[测试版] AMD GPU 二进制文件现已可用
从 PyTorch 1.8 开始,我们增加了对 ROCm wheel 的支持,为使用 AMD GPU 提供了简便的入门。您只需访问标准 PyTorch 安装选择器,选择 ROCm 作为安装选项并执行提供的命令。
感谢您的阅读,如果您对这些更新感到兴奋并希望参与 PyTorch 的未来,我们鼓励您加入讨论论坛并提出 GitHub 问题。
干杯!
PyTorch 团队