我们很高兴地宣布 PyTorch 1.8 的发布。自 1.7 版本以来,此版本包含了超过 3,000 次提交。它包含了编译、代码优化、科学计算前端 API 以及通过 pytorch.org 提供的二进制文件支持 AMD ROCm 的重大更新和新功能。它还为管道和模型并行化以及梯度压缩提供了改进的大规模训练功能。其中一些亮点包括:
- 支持通过
torch.fx
进行 Python 到 Python 的函数式转换; - 添加或稳定了支持 FFTs (
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 风格的 FFTs
作为 PyTorch 支持科学计算目标的一部分,我们投入了改进 FFT 支持,并在 PyTorch 1.8 中发布了 torch.fft
模块。此模块实现了与 NumPy 的 np.fft
模块相同的功能,但支持硬件加速和自动微分。
[测试版] 通过 torch.linalg
支持 NumPy 风格的线性代数函数
torch.linalg
模块模仿 NumPy 的 np.linalg 模块,为常见的线性代数操作提供了 NumPy 风格的支持,包括 Cholesky 分解、行列式、特征值等等。
[测试版] 使用 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 中的 vanilla allreduce 来控制如何跨 workers 传输梯度。提供了几个内置的通信钩子,包括 PowerSGD,用户可以轻松应用其中任何一个钩子来优化通信。此外,通信钩子接口还可以支持用户自定义的通信策略,以实现更高级的用例。
分布式训练的附加原型功能
除了此版本中主要的稳定版和测试版分布式训练功能外,我们的 nightly 构建中还提供了许多原型功能,供您试用并提供反馈。我们已在下方链接了草稿文档以供参考。
- (原型)ZeroRedundancyOptimizer – 基于并与 Microsoft 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 Mobile
PyTorch Mobile 的支持正在扩展,提供了一套新的教程,以帮助新用户更快地在设备上启动模型,并为现有用户提供一个工具,以从我们的框架中获得更多收益。这些包括:
我们新的演示应用程序还包括图像分割、目标检测、神经机器翻译、问答和视觉转换器的示例。它们在 iOS 和 Android 上均可用。
除了 MobileNetV3 和其他模型的 CPU 性能改进之外,我们还改进了 Android GPU 后端原型,以实现更广泛的模型覆盖和更快的推理。
最后,我们在此版本中发布 PyTorch Mobile Lite Interpreter 作为原型功能。Lite Interpreter 允许用户减少运行时二进制文件大小。请尝试这些功能并在 PyTorch 论坛上向我们发送您的反馈。所有最新更新都可以在 PyTorch Mobile 页面上找到。
[原型] PyTorch Mobile 精简解释器
PyTorch 精简解释器是 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。它通过添加对函数式编程的支持和自动化量化过程来改进 Eager 模式量化,尽管人们可能需要重构模型以使其与 FX 图模式量化(可使用 torch.fx
进行符号追踪)兼容。
硬件支持
[测试版] 在 C++ 中为新后端扩展 PyTorch 调度器的能力
在 PyTorch 1.8 中,您现在可以创建在 pytorch/pytorch
仓库之外的新的树外设备。下面链接的教程展示了如何注册您的设备并使其与原生 PyTorch 设备保持同步。
[测试版] AMD GPU 二进制文件现已可用
从 PyTorch 1.8 开始,我们增加了对 ROCm wheel 的支持,为使用 AMD GPU 提供了便捷的入门体验。您只需访问标准的 PyTorch 安装选择器,选择 ROCm 作为安装选项,然后执行提供的命令即可。
感谢您的阅读,如果您对这些更新感到兴奋并希望参与 PyTorch 的未来,我们鼓励您加入讨论论坛并提出 GitHub 问题。
干杯!
PyTorch 团队