由 PyTorch 团队

我们非常高兴地宣布 PyTorch 1.9 的发布。自 1.8 版本以来,共有 398 位贡献者提交了超过 3,400 次提交。发布说明在此处可用。亮点包括:

  1. 对科学计算的支持进行了重大改进,包括 torch.linalg、torch.special 和复数自动微分
  2. 在 Mobile Interpreter 中实现了对设备上二进制大小的重大改进
  3. 通过将 TorchElastic 集成到 PyTorch 核心中,原生支持弹性容错训练
  4. PyTorch RPC 框架的重大更新,以支持具有 GPU 支持的规模化分布式训练
  5. 新的 API 以优化模型推理部署的性能和打包
  6. 支持 PyTorch Profiler 中的分布式训练、GPU 利用率和 SM 效率

除了 1.9 版本外,我们还发布了 PyTorch 库的重大更新,您可以在本博客文章中了解相关信息。

我们想感谢社区对本次最新版本的支持和工作。我们特别感谢 Quansight 和 Microsoft 的贡献。

PyTorch 版本中的功能分为稳定版、测试版和原型版。您可以在本博客文章中了解更多关于定义的信息。

前端 API

(稳定版) torch.linalg

在 1.9 版本中,torch.linalg 模块将转为稳定版发布。线性代数对于深度学习和科学计算至关重要,torch.linalg 模块通过实现 NumPy 线性代数模块(现在支持加速器和自动微分)中的每个函数以及更多功能(如 torch.linalg.matrix_norm 和 torch.linalg.householder_product),扩展了 PyTorch 对线性代数的支持。这使得该模块对使用过 NumPy 的用户来说非常熟悉。请参阅此处文档。

我们计划下周发布另一篇博客文章,详细介绍 torch.linalg 模块的更多细节!

(稳定) 复杂自动微分

复杂自动微分功能,在 PyTorch 1.8 中作为 beta 版发布,现已稳定。自 beta 版发布以来,我们已将 PyTorch 1.9 中超过 98%的操作符支持复杂自动微分,通过添加更多 OpInfos 来改进复杂操作符的测试,并通过 TorchAudio 迁移到原生复数张量增加了更多的验证(参阅此问题)。

此功能为用户提供计算复数梯度和使用复数变量优化实值损失函数的功能。这是 PyTorch 中多个当前和未来潜在用户(如 TorchAudio、ESPNet、Asteroid 和 FastMRI)所必需的功能。有关详细信息,请参阅文档。

(稳定) torch.use_deterministic_algorithms()

为了帮助调试和编写可重复的程序,PyTorch 1.9 包含了一个 torch.use_determinstic_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(accum=False)。有关更多详细信息,请参阅文档和可重复性说明。

(Beta) torch.special

现在在测试版中,torch.special 模块类似于 SciPy 的特殊模块,现在可用。此模块包含许多用于科学计算和与分布(如 iv、ive、erfcx、logerfc 和 logerfcx)一起工作的函数。有关更多详细信息,请参阅文档。

(Beta) nn.Module 参数化

nn.Module 参数化允许用户对 nn.Module 的任何参数或缓冲区进行参数化,而无需修改 nn.Module 本身。它允许您在不使用特殊优化方法的情况下约束参数存在的空间。

这还包含了对 PyTorch 1.9 的 spectral_norm 参数化的一种新实现。更多参数化(weight_norm、矩阵约束和部分剪枝)将添加到该功能中,以便在 1.10 版本中使该功能稳定。有关更多详细信息,请参阅文档和教程。

PyTorch 移动

(Beta) 移动解释器

我们正在发布移动翻译器,这是 PyTorch 运行时的精简版,目前处于测试阶段。翻译器将在边缘设备上执行 PyTorch 程序,并减少二进制文件大小。

移动翻译器是 PyTorch Mobile 最受欢迎的功能之一。与当前设备上的运行时相比,这个新版本将显著减少二进制文件大小。为了您能够通过我们的翻译器(可以将二进制文件大小减少高达 75%)获得二进制文件大小的改进,请按照以下说明操作。例如,使用移动翻译器,我们可以将 MobileNetV2 在 arm64-v7a Android 上压缩到 2.6 MB。在这个最新版本中,我们通过提供预构建的 iOS 和 Android 库,使集成翻译器变得更加简单。

TorchVision 库

从 1.9 版本开始,用户可以在 iOS/Android 应用中使用 TorchVision 库。Torchvision 库包含 C++ TorchVision 操作,需要与 iOS 的主 PyTorch 库链接,对于 Android,可以作为 gradle 依赖项添加。这允许使用预构建的 TorchVision MaskRCNN 操作符进行目标检测和分割。要了解更多关于该库的信息,请参阅我们的教程和演示应用。

演示应用

我们将发布一个基于 PyTorch Video 库的新视频应用和一个基于最新 torchaudio、wave2vec 模型的更新版语音识别应用。这两个应用都可在 iOS 和 Android 上使用。此外,我们还更新了七个计算机视觉和三个自然语言处理演示应用,包括 HuggingFace DistilBERT 和 DeiT 视觉 Transformer 模型,使用 PyTorch Mobile v1.9。随着这两个应用的加入,我们现在提供了一套完整的覆盖图像、文本、音频和视频的演示应用。要开始使用,请查看我们的 iOS 演示应用和 Android 演示应用。

分布式训练

(Beta)TorchElastic 现在是核心部分

TorchElastic,一年前在 pytorch/elastic GitHub 仓库开源,是 PyTorch 工作进程的运行器和协调器。从那时起,它已被各种分布式 torch 用例采用:1)deepspeech.pytorch 2)pytorch-lightning 3)Kubernetes CRD。现在,它已成为 PyTorch 核心的一部分。

如其名所示,TorcheElastic 的核心功能是优雅地处理扩展事件。弹性的一个显著推论是,对等发现和排名分配已内置到 TorchElastic 中,使用户能够在不需要集群调度器的情况下在可抢占实例上运行分布式训练。顺便提一下,etcd 曾经是 TorchElastic 的硬依赖。随着上游的更新,这种情况已经不再存在,因为我们已经添加了一个基于 c10d::Store 的“独立” rendezvous。有关更多详细信息,请参阅文档。

(Beta)分布式训练更新

除了 TorchElastic 之外,分布式包中还有许多可用的 beta 功能:

  • (Beta) CUDA 支持在 RPC 中可用:与 CPU RPC 和通用 RPC 框架相比,CUDA RPC 是 P2P Tensor 通信的一种更高效的方式。它建立在 TensorPipe 之上,可以根据 Tensor 设备类型和调用方及被调用方上的通道可用性自动选择每个 Tensor 的通信通道。现有的 TensorPipe 通道包括 NVLink、InfiniBand、SHM、CMA、TCP 等。请参阅此菜谱了解 CUDA RPC 如何帮助实现与 CPU RPC 相比的 34 倍速度提升。

  • (Beta) ZeroRedundancyOptimizer:ZeroRedundancyOptimizer 可以与 DistributedDataParallel 结合使用,以减少每个进程的优化器状态的大小。ZeroRedundancyOptimizer 的灵感来自 DeepSpeed/ZeRO 项目和 Marian,其中每个进程的优化器拥有模型参数的片段及其相应的优化器状态。当运行 step() 时,每个优化器只更新自己的参数,然后使用集体通信同步所有进程的更新参数。请参阅此文档和此教程了解详细信息。

  • (Beta) 支持分布式集体分析:PyTorch 的剖析工具 torch.profiler 和 torch.autograd.profiler 能够分析分布式集体和点对点通信原语,包括 allreduce、alltoall、allgather、send/recv 等。这适用于 PyTorch 原生支持的所有后端:gloo、mpi 和 nccl。这可以用于调试性能问题,分析包含分布式通信的跟踪,并深入了解使用分布式训练的应用程序的性能。要了解更多信息,请参阅此文档。

性能优化和工具

(Stable) API 冻结

模块冻结是将模块参数和属性值内联为常量到 TorchScript 内部表示的过程。这允许进一步优化和专门化您的程序,包括 TorchScript 优化和降低到其他后端。它被 optimize_for_mobile API、ONNX 等使用。

推荐将模型部署时进行冻结。这有助于 TorchScript JIT 优化移除训练、调整或调试 PyTorch 模型所需的额外开销和记录。它启用了在非冻结图上不语义有效的图融合,例如融合卷积-归一化。有关更多详细信息,请参阅文档。

(Beta) PyTorch 性能分析器

新的 PyTorch 分析器已升级至 Beta 版本,并利用 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()

更多使用示例可以在分析器食谱页面上找到。

PyTorch 分析器 Tensorboard 插件新增以下功能:

  • NCCL 的分布式训练摘要视图和通信概览
  • 跟踪视图和 GPU 操作视图中的 GPU 利用率及 SM 效率
  • 内存分析视图
  • 从 Microsoft VSCode 启动时跳转到源代码
  • 能够从云对象存储系统中加载跟踪

(Beta) 推理模式 API

推理模式 API 允许在保持安全并确保不会计算任何错误梯度的情况下,显著提高推理工作负载的速度。当不需要自动微分时,它提供最佳性能。有关详细信息,请参阅推理模式的文档以及解释何时使用它以及与 no_grad 模式差异的文档。

(Beta) 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 问题。要获取 PyTorch 的最新消息,请关注我们的 Facebook、Twitter、Medium、YouTube 或 LinkedIn。

喝彩!

PyTorch 团队