我们很高兴地宣布 PyTorch 1.9 正式发布。自 1.8 版本以来,该版本由 398 位贡献者提交了超过 3,400 次代码合并。发布说明可在 此处查看。主要亮点包括:
- 科学计算支持方面的重大改进,包括 torch.linalg、torch.special 和复数自动求导 (Complex Autograd)
- 通过 Mobile Interpreter 显著缩小了端侧二进制文件体积
- 将 TorchElastic 合并至 PyTorch 核心,原生支持弹性容错训练
- 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 模块扩展了 PyTorch 在此方面的支持,实现了 NumPy 线性代数模块中的所有函数(现已支持加速器和自动求导),并添加了更多功能,如 torch.linalg.matrix_norm 和 torch.linalg.householder_product。这使得熟悉 NumPy 的用户可以无缝上手。请参阅官方文档。
我们计划下周发布另一篇关于 torch.linalg 模块的详细博客文章!
(稳定版) 复数自动求导 (Complex Autograd)
在 PyTorch 1.8 中作为测试版发布的复数自动求导功能现已稳定。自测试版发布以来,我们已将复数自动求导的支持扩展到 PyTorch 1.9 中超过 98% 的算子,通过增加更多 OpInfos 改进了对复数算子的测试,并通过 TorchAudio 迁移至原生复数张量增强了验证(请参阅此议题)。
该功能为用户提供了计算复数梯度以及使用复数变量优化实值损失函数的能力。这是 TorchAudio、ESPNet、Asteroid 和 FastMRI 等目前及未来下游使用 PyTorch 复数功能的项目所需的核心特性。更多详情请参阅文档。
(稳定版) 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 的 special 模块,torch.special 模块现已作为测试版提供。该模块包含了许多对科学计算和处理分布有用的函数,如 iv、ive、erfcx、logerfc 和 logerfcx。更多详情请参阅官方文档。
(测试版) nn.Module 参数化 (Parameterization)
nn.Module 参数化允许用户在不修改 nn.Module 本身的情况下,对其任意参数或缓冲区进行参数化。它允许你在不需要特殊优化方法的情况下,约束参数所在的取值空间。
这还包含了 PyTorch 1.9 中 spectral_norm 参数化的新实现。未来将为此功能添加更多参数化方法(weight_norm、矩阵约束和部分剪枝),以期在 1.10 版本达到稳定。更多详情请参阅文档和教程。
PyTorch Mobile
(测试版) Mobile Interpreter
我们发布了 Mobile Interpreter,这是一个精简版的 PyTorch 运行时,目前处于测试阶段。Interpreter 将在边缘设备上执行 PyTorch 程序,并显著减少二进制文件大小。
Mobile Interpreter 是 PyTorch Mobile 最受期待的功能之一。与当前的端侧运行时相比,此版本将显著减小二进制体积。若要利用解释器实现二进制大小优化(对于典型应用程序,体积最高可减少约 75%),请按照这些说明操作。例如,使用 Mobile Interpreter,我们在 arm64-v7a Android 上压缩后的 MobileNetV2 大小可达到 2.6 MB。随着最新版本的发布,我们通过提供 iOS 和 Android 的预构建库,简化了集成过程。
TorchVision 库
从 1.9 版本开始,用户可以在其 iOS/Android 应用中使用 TorchVision 库。TorchVision 库包含 C++ TorchVision 算子,在 iOS 上需要与主 PyTorch 库链接,在 Android 上则可以作为 gradle 依赖项添加。这使得使用预构建的 MaskRCNN 算子进行目标检测和分割成为可能。要了解更多信息,请参考我们的教程和演示应用。
演示应用
我们发布了一个基于 PyTorch Video 库的新视频应用,以及一个基于最新 torchaudio 和 wave2vec 模型的更新版语音识别应用。两者均可在 iOS 和 Android 上获取。此外,我们已将七个计算机视觉演示应用和三个自然语言处理演示应用(包括 HuggingFace DistilBERT 和 DeiT 视觉 Transformer 模型)更新为 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 Scheduler 即可在可抢占实例上运行分布式训练。顺便提一下,etcd 曾经是 TorchElastic 的硬依赖,随着此次上游集成,情况已不再如此,因为我们添加了一个基于 c10d::Store 的“独立”rendezvous。更多详情请参阅文档。
(测试版) 分布式训练更新
除了 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。这可用于调试性能问题、分析包含分布式通信的追踪记录,并深入了解使用分布式训练的应用程序性能。要了解更多信息,请参阅此文档。
性能优化和工具
(稳定版) 冻结 (Freezing) API
模块冻结是将模块参数和属性值作为常量内联到 TorchScript 内部表示中的过程。这允许对程序进行进一步的优化和特殊化,无论是针对 TorchScript 优化还是降低到其他后端。它被 optimize_for_mobile API、ONNX 等所使用。
推荐在模型部署时使用冻结。它有助于 TorchScript JIT 优化器消除训练、微调或调试 PyTorch 模型所需的开销和簿记工作。它启用了在非冻结图中语义上不成立的图融合——例如 Conv-BN 融合。更多详情,请参阅文档。
(测试版) PyTorch Profiler

全新的 PyTorch Profiler 正式进入测试阶段,利用 Kineto 进行 GPU 分析,使用 TensorBoard 进行可视化,现已成为我们教程和文档中的标准分析工具。
PyTorch 1.9 将新的 torch.profiler API 支持范围扩大到了更多构建版本,包括 Windows 和 Mac。在大多数情况下,推荐使用此 API 而非之前的 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()
更多使用示例可以在 Profiler 示例页面找到。
PyTorch Profiler Tensorboard 插件的新功能包括:
- 带有 NCCL 通信概览的分布式训练摘要视图
- 追踪视图和 GPU 算子视图中的 GPU 利用率和 SM 效率
- 内存分析视图
- 从 Microsoft VSCode 启动时跳转至源码
- 能够从云对象存储系统加载追踪记录
(测试版) 推理模式 (Inference Mode) API
推理模式 API 可为推理负载带来显著加速,同时保持安全性并确保绝不会计算出错误的梯度。在不需要自动求导时,它提供了最佳的性能。更多详情,请参阅推理模式文档以及说明何时使用它以及它与 no_grad 模式的区别的文档。
(测试版) torch.package
torch.package 是一种以自包含、稳定格式打包 PyTorch 模型的新方法。包将包含模型的配置数据(如参数、缓冲区)及其代码(模型架构)。将模型及其全套 Python 依赖项(结合带有固定版本的 conda 环境描述)打包在一起,可以轻松复现训练过程。将模型表示为一个自包含的工件,也允许它在整个生产 ML 流水线中发布和传输,同时保留纯 Python 表示的灵活性。更多详情请参阅文档。
(原型) prepare_for_inference
prepare_for_inference 是一项新的原型功能,它接收一个模块并根据设备执行图层面的优化,以提高推理性能。它旨在成为一种 PyTorch 原生选项,对用户的开发流程只需最小的改动。更多详情,请查看 TorchScript 版本文档 或 FX 版本文档。
(原型) TorchScript 中的 Profile-directed typing
TorchScript 对源代码有硬性要求,必须包含类型注解才能编译成功。长期以来,仅能通过试错法(即逐一修复 torch.jit.script 生成的类型检查错误)来添加缺失或不正确的类型注解,过程低效且耗时。现在,我们通过利用 MonkeyType 等现有工具,启用了针对 torch.jit.script 的 Profile-directed typing,使得该过程变得更加简单、快捷和高效。更多详情请参阅文档。
感谢阅读。如果您对这些更新感兴趣并想加入 PyTorch 社区,我们鼓励您加入讨论论坛并在 GitHub 上提交议题。要获取 PyTorch 的最新资讯,请关注我们的 Facebook、Twitter、Medium、YouTube 或 LinkedIn。
干杯!
PyTorch 团队