我们很高兴地宣布 PyTorch 1.9 发布。自 1.8 版本以来,此版本包含了由 398 位贡献者提交的 3,400 多次提交。发布说明可在此处获取: 此处。主要亮点包括:
- 科学计算支持方面的重大改进,包括 torch.linalg、torch.special 和复杂自动微分(Complex Autograd)
- 通过移动端解释器(Mobile Interpreter)大幅减小设备上二进制文件大小
- 通过将 TorchElastic 上游集成到 PyTorch Core 中,原生支持弹性容错训练
- PyTorch RPC 框架的重大更新,支持带 GPU 的大规模分布式训练
- 新的 API,用于优化模型推理部署的性能和打包
- PyTorch 分析器(Profiler)中支持分布式训练、GPU 利用率和 SM 效率
除 1.9 版本外,我们还发布了 PyTorch 库的重大更新,您可以在 这篇博客文章中阅读相关内容。
我们要感谢社区对本次最新发布的支持和贡献。我们特别感谢 Quansight 和 Microsoft 的贡献。
PyTorch 版本中的功能分为稳定版 (Stable)、测试版 (Beta) 和原型版 (Prototype)。您可以在这篇博客文章中了解更多定义。
前端 API
(稳定版) torch.linalg
在 1.9 版本中,torch.linalg 模块将进入稳定版发布。线性代数对于深度学习和科学计算至关重要,torch.linalg 模块通过实现 NumPy 线性代数模块中的所有函数(现在支持加速器和自动微分)以及更多功能,如 torch.linalg.matrix_norm 和 torch.linalg.householder_product,扩展了 PyTorch 对线性代数的支持。这使得该模块对使用过 NumPy 的用户来说非常熟悉。请参阅 此处的文档。
我们计划下周发布另一篇关于 torch.linalg 模块的更详细博客文章!
(稳定版) 复杂自动微分(Complex Autograd)
在 PyTorch 1.8 中作为测试版发布的复杂自动微分功能现已稳定。自测试版发布以来,我们已将 PyTorch 1.9 中对复杂自动微分的支持扩展到 98% 以上的运算符,通过添加更多 OpInfos 改进了复杂运算符的测试,并通过 TorchAudio 迁移到原生复杂张量增加了验证(参阅 此问题)。
此功能为用户提供了计算复杂梯度和使用复杂变量优化实值损失函数的功能。对于 PyTorch 中复杂数字的许多现有和下游潜在用户,如 TorchAudio、ESPNet、Asteroid 和 FastMRI,这是一个必要的功能。有关更多详细信息,请参阅 文档。
(稳定版) torch.use_deterministic_algorithms()
为帮助调试和编写可复现程序,PyTorch 1.9 包含一个 torch.use_deterministic_algorithms 选项。启用此设置后,操作(如果可能)将确定性地执行,或者如果它们可能非确定性地执行,则会抛出运行时错误。以下是一些示例:
>>> a = torch.randn(100, 100, 100, device='cuda').to_sparse()
>>> b = torch.randn(100, 100, 100, device='cuda')
# Sparse-dense CUDA bmm is usually nondeterministic
>>> torch.bmm(a, b).eq(torch.bmm(a, b)).all().item()
False
>>> torch.use_deterministic_algorithms(True)
# Now torch.bmm gives the same result each time, but with reduced performance
>>> torch.bmm(a, b).eq(torch.bmm(a, b)).all().item()
True
# CUDA kthvalue has no deterministic algorithm, so it throws a runtime error
>>> torch.zeros(10000, device='cuda').kthvalue(1)
RuntimeError: kthvalue CUDA does not have a deterministic implementation...
PyTorch 1.9 还为许多索引操作添加了确定性实现,包括 index_add、index_copy 和 index_put with accum=False。有关更多详细信息,请参阅 文档 和 可复现性说明。
(测试版) torch.special
一个类似于 SciPy 的特殊模块的 torch.special 模块现已推出测试版。此模块包含许多对科学计算和处理分布有用的函数,例如 iv、ive、erfcx、logerfc 和 logerfcx。有关更多详细信息,请参阅 文档。
(测试版) nn.Module 参数化
nn.Module
参数化允许用户对 nn.Module
的任何参数或缓冲区进行参数化,而无需修改 nn.Module
本身。它允许您限制参数所处的空间,而无需特殊的优化方法。
这还包含 PyTorch 1.9 中 spectral_norm
参数化的新实现。此功能(权重范数、矩阵约束和剪枝的一部分)将添加更多参数化,以便该功能在 1.10 中变得稳定。有关更多详细信息,请参阅 文档 和 教程。
PyTorch Mobile
(测试版) 移动端解释器(Mobile Interpreter)
我们正在发布移动端解释器(Mobile Interpreter)的测试版,它是 PyTorch 运行时的精简版本。该解释器将在边缘设备上执行 PyTorch 程序,并减小二进制文件的大小。
移动端解释器是 PyTorch Mobile 最受欢迎的功能之一。这个新版本将与当前的设备上运行时相比,显著减小二进制文件的大小。为了通过我们的解释器获得二进制文件大小的改进(对于典型应用程序,可以将二进制文件大小减少多达约 75%),请遵循这些说明。例如,使用移动端解释器,我们可以在 arm64-v7a Android 上使用 MobileNetV2 达到 2.6 MB 的压缩大小。通过本次最新发布,我们通过为 iOS 和 Android 提供预构建库,使集成解释器变得更加简单。
TorchVision 库
从 1.9 版本开始,用户可以在其 iOS/Android 应用上使用 TorchVision 库。TorchVision 库包含 C++ TorchVision ops,需要与主要的 PyTorch 库一起链接到 iOS,对于 Android,可以将其添加为 Gradle 依赖项。这允许使用 TorchVision 预构建的 MaskRCNN 运算符进行对象检测和分割。要了解有关该库的更多信息,请参阅我们的教程和 演示应用程序。
演示应用程序
我们正在发布一个基于 PyTorch Video 库的新视频应用程序,以及一个基于最新 torchaudio, wave2vec 模型更新的语音识别应用程序。两者都可在 iOS 和 Android 上使用。此外,我们已将七个计算机视觉和三个自然语言处理演示应用程序(包括 HuggingFace DistilBERT 和 DeiT 视觉转换器模型)更新到 PyTorch Mobile v1.9。随着这两个应用程序的加入,我们现在提供了一整套涵盖图像、文本、音频和视频的演示应用程序。要开始使用,请查看我们的 iOS 演示应用程序 和 Android 演示应用程序。

分布式训练
(测试版) TorchElastic 现已成为核心部分
TorchElastic,一年前在 pytorch/elastic GitHub 仓库中开源,是 PyTorch 工作进程的运行器和协调器。自那时以来,它已被各种分布式 Torch 用例采用:1) deepspeech.pytorch 2) pytorch-lightning 3) Kubernetes CRD。现在,它已成为 PyTorch 核心的一部分。
顾名思义,TorchElastic 的核心功能是优雅地处理扩展事件。弹性的一个显著推论是,TorchElastic 内置了对等发现和排名分配,使用户无需 Gang 调度程序即可在可抢占实例上运行分布式训练。顺便提一下,etcd 曾是 TorchElastic 的硬依赖项。通过上游集成,情况不再如此,因为我们已经添加了一个基于 c10d::Store 的“独立”会合点。有关更多详细信息,请参阅 文档。
(测试版) 分布式训练更新
除了 TorchElastic,分布式包中还有许多测试版功能可用
- (测试版) RPC 中提供 CUDA 支持:与 CPU RPC 和通用 RPC 框架相比,CUDA RPC 是一种更高效的 P2P 张量通信方式。它构建在 TensorPipe 之上,可以根据张量设备类型以及调用者和被调用者上的通道可用性,自动为每个张量选择一个通信通道。现有 TensorPipe 通道涵盖 NVLink、InfiniBand、SHM、CMA、TCP 等。请参阅 此教程,了解 CUDA RPC 如何实现比 CPU RPC 快 34 倍的加速。
- (测试版) ZeroRedundancyOptimizer:ZeroRedundancyOptimizer 可以与 DistributedDataParallel 结合使用,以减少每个进程的优化器状态大小。ZeroRedundancyOptimizer 的思想来源于 DeepSpeed/ZeRO 项目 和 Marian,其中每个进程中的优化器拥有模型参数及其相应优化器状态的一部分。在运行
step()
时,每个优化器只更新自己的参数,然后使用集体通信同步所有进程中更新的参数。请参阅 此文档 和 此教程 以了解更多信息。 - (测试版) 支持分布式集合操作分析:PyTorch 的分析器工具 torch.profiler 和 torch.autograd.profiler 能够分析分布式集合操作和点对点通信原语,包括 allreduce、alltoall、allgather、send/recv 等。这已为 PyTorch 原生支持的所有后端启用:gloo、mpi 和 nccl。这可用于调试性能问题、分析包含分布式通信的跟踪,并深入了解使用分布式训练的应用程序的性能。要了解更多信息,请参阅 此文档。
性能优化和工具
(稳定版) 冻结 API
模块冻结是将模块参数和属性值作为常量内联到 TorchScript 内部表示的过程。这允许对程序进行进一步优化和专门化,无论是 TorchScript 优化还是降级到其他后端。它被 optimize_for_mobile API、ONNX 和其他工具使用。
建议将冻结用于模型部署。它有助于 TorchScript JIT 优化消除训练、调优或调试 PyTorch 模型所需的开销和簿记。它支持在非冻结图上语义无效的图融合——例如融合 Conv-BN。有关更多详细信息,请参阅 文档。
(测试版) PyTorch 分析器

新的 PyTorch 分析器升级到测试版,它利用 Kineto 进行 GPU 分析,利用 TensorBoard 进行可视化,现已成为我们教程和文档中的标准。
PyTorch 1.9 将新的 torch.profiler API 的支持扩展到更多构建版本,包括 Windows 和 Mac,并且在大多数情况下推荐使用它来代替之前的 torch.autograd.profiler API。新的 API 支持现有的分析器功能,与 CUPTI 库(仅限 Linux)集成以跟踪设备上的 CUDA 内核,并支持长时间运行的任务,例如:
def trace_handler(p):
output = p.key_averages().table(sort_by="self_cuda_time_total", row_limit=10)
print(output)
p.export_chrome_trace("/tmp/trace_" + str(p.step_num) + ".json")
with profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
# schedule argument specifies the iterations on which the profiler is active
schedule=torch.profiler.schedule(
wait=1,
warmup=1,
active=2),
# on_trace_ready argument specifies the handler for the traces
on_trace_ready=trace_handler
) as p:
for idx in range(8):
model(inputs)
# profiler will trace iterations 2 and 3, and then 6 and 7 (counting from zero)
p.step()
更多使用示例可在 分析器教程页面上找到。
PyTorch Profiler Tensorboard 插件具有以下新功能:
- 分布式训练摘要视图,包含 NCCL 通信概览
- 跟踪视图和 GPU 运算符视图中的 GPU 利用率和 SM 效率
- 内存分析视图
- 从 Microsoft VSCode 启动时跳转到源代码
- 从云对象存储系统加载跟踪的能力
(测试版) 推理模式 API
推理模式 API 可以在推理工作负载中显著加速,同时保持安全并确保不会计算出不正确的梯度。当不需要自动微分时,它提供最佳性能。有关更多详细信息,请参阅 推理模式本身的文档 以及 解释何时使用它以及与 no_grad 模式区别的文档。
(测试版) torch.package
torch.package 是一种以自包含、稳定格式打包 PyTorch 模型的新方法。一个包将包含模型的数据(例如参数、缓冲区)及其代码(模型架构)。将模型及其完整的 Python 依赖项打包,并结合带有固定版本的 conda 环境描述,可以轻松重现训练。以自包含的工件形式表示模型还可以使其在生产 ML 管道中发布和传输,同时保留纯 Python 表示的灵活性。有关更多详细信息,请参阅 文档。
(原型) prepare_for_inference
prepare_for_inference 是一项新的原型功能,它接收一个模块并根据设备执行图级优化以提高推理性能。它旨在成为 PyTorch 原生选项,对用户工作流程的更改最小。有关更多详细信息,请参阅 Torchscript 版本的 此处 文档或 FX 版本的 此处 文档。
(原型) TorchScript 中的配置文件导向类型推断
TorchScript 对源代码有严格要求,需要带有类型注解才能成功编译。长期以来,只能通过试错法(即通过逐一修复 torch.jit.script 生成的类型检查错误)来添加缺失或不正确的类型注解,这效率低下且耗时。现在,我们通过利用 MonkeyType 等现有工具,为 torch.jit.script 启用了配置文件导向类型推断,这使得该过程更加简单、快速、高效。有关更多详细信息,请参阅 文档。
感谢阅读。如果您对这些更新感兴趣并希望加入 PyTorch 社区,我们鼓励您加入 讨论论坛 并 提交 GitHub issue。要获取 PyTorch 的最新消息,请关注我们的 Facebook、Twitter、Medium、YouTube 或 LinkedIn。
干杯!
PyTorch 团队