博客

PyTorch 1.10 发布，包括 CUDA Graphs API、前端和编译器改进

作者： 2021年10月21日2024年11月15日暂无评论

我们非常高兴地宣布 PyTorch 1.10 正式发布。自 1.9 版本以来，共有 426 位贡献者提交了超过 3,400 次代码更改。我们要衷心感谢社区为持续改进 PyTorch 所做的贡献。

PyTorch 1.10 的更新重点在于改善训练过程、提升 PyTorch 性能以及增强开发者易用性。完整发行说明请见 此处。主要亮点包括：

  1. 集成了 CUDA Graphs API，以降低 CUDA 工作负载的 CPU 开销。
  2. 多个前端 API（如 FX、torch.special 和 nn.Module Parametrization）已从测试版转为稳定版。
  3. JIT 编译器中自动融合的支持范围从 GPU 扩展到了 CPU。
  4. Android NNAPI 支持现已进入测试阶段。

伴随 1.10 版本的发布，我们还对 PyTorch 库进行了重大更新，您可以阅读 这篇博文 了解详情。

前端 API

(稳定版) 基于 FX 的 Python 代码转换

FX 为转换和降低（lowering）PyTorch 程序提供了一个 Python 风格的平台。它是一个供转换编写者使用的工具包，旨在促进函数和 nn.Module 实例的 Python 到 Python 转换。该工具包旨在支持 Python 语言语义的一个子集——而非完整语言——从而简化转换的实现。在 1.10 版本中，FX 正式进入稳定阶段。

您可以通过 官方文档 以及使用 torch.fx 实现程序转换的 GitHub 示例 了解更多关于 FX 的信息。

(稳定版) torch.special

一个类似于 SciPy 特殊函数模块 的 torch.special 模块现已稳定可用。该模块包含 30 种运算，包括伽马函数、贝塞尔函数和（高斯）误差函数等。

请参阅此 文档 以获取更多详细信息。

(稳定版) nn.Module 参数化 (Parametrization)

nn.Module 参数化功能现已稳定。该功能允许用户在不修改 nn.Module 本身的情况下，对其任意参数或缓冲区进行参数化。此版本增加了权重归一化（weight_norm）、正交参数化（矩阵约束和部分剪枝功能），并为您创建自定义参数化提供了更大的灵活性。

请参阅此 教程 和通用 文档 以获取更多详细信息。

(测试版) CUDA Graphs API 集成

PyTorch 现已集成 CUDA Graphs API，旨在降低 CUDA 工作负载的 CPU 开销。

CUDA Graphs 极大降低了 CPU 密集型 CUDA 工作负载的 CPU 开销，从而通过提高 GPU 利用率来提升性能。对于分布式工作负载，CUDA Graphs 还能减少抖动；由于并行工作负载必须等待最慢的进程，减少抖动可提高整体并行效率。

集成允许在由 CUDA Graphs 捕获的网络部分与因图限制而无法捕获的部分之间实现无缝互操作。

请阅读此 说明 以获取更多详细信息和示例，并参阅通用 文档 获取补充信息。

[测试版] 共轭视图 (Conjugate View)

PyTorch 对复数张量的共轭运算 (torch.conj()) 现在是一个常数时间操作，它返回输入张量的一个视图，并设置一个共轭位（可通过调用 torch.is_conj() 查看）。这一点已被矩阵乘法、点积等其他多种 PyTorch 运算所利用，实现了将共轭操作与运算融合，从而在 CPU 和 CUDA 上均带来了显著的性能提升和内存节省。

分布式训练

分布式训练发布现已转为稳定版

在 1.10 版本中，分布式包中有多个功能从测试版转为稳定版：

  • (稳定版) 远程模块 (Remote Module)：该功能允许用户像使用本地模块一样操作远程 worker 上的模块，且 RPC 对用户是透明的。请参阅此 文档 了解更多详情。
  • (稳定版) DDP 通信钩子 (DDP Communication Hook)：该功能允许用户自定义 DDP 在各进程间同步梯度的方式。请参阅此 文档 了解更多详情。
  • (稳定版) ZeroRedundancyOptimizer：该功能可与 DistributedDataParallel 结合使用，以减少每个进程的优化器状态大小。随着本次稳定发布，它现在可以处理发送到不同数据并行 worker 的不均匀输入。查看此 教程。我们还改进了参数分区算法，以更好地平衡各进程间的内存和计算开销。请参阅此 文档 和此 教程 了解更多信息。

性能优化和工具

[测试版] TorchScript 中的配置导向类型推断 (Profile-directed typing)

TorchScript 强制要求源代码必须具有类型注解才能编译成功。长期以来，开发者只能通过反复试验来添加缺失或不正确的类型注解（即逐一修复 torch.jit.script 生成的类型检查错误），这既低效又耗时。

现在，我们通过利用 MonkeyType 等现有工具，为 torch.jit.script 启用了配置导向类型推断，使得这一过程变得更加简便、快速且高效。更多详情请参阅 文档。

(测试版) CPU 融合 (CPU Fusion)

在 PyTorch 1.10 中，我们为 CPU 添加了一个基于 LLVM 的 JIT 编译器，它可以将一系列 torch 库调用融合在一起以提高性能。尽管我们在 GPU 上已经拥有此功能一段时间了，但这是我们首次将编译功能带到 CPU。
您可以在此 Colab 笔记本 中亲自查看一些性能评估结果。

(测试版) PyTorch Profiler

PyTorch Profiler 的目标是定位执行步骤中最耗时和/或最耗内存的部分，并可视化 GPU 和 CPU 之间的工作负载分布。PyTorch 1.10 包含以下关键功能：

  • 增强的内存视图 (Enhanced Memory View)：这能帮助您更好地理解内存使用情况。该工具通过显示程序运行过程中各个时间点的活跃内存分配情况，帮助您避免内存溢出 (OOM) 错误。
  • 增强的自动建议 (Enhanced Automated Recommendations)：这有助于提供自动化的性能建议，以优化您的模型。这些工具会针对批大小 (batch size)、TensorCore、内存缩减技术等提出修改建议。
  • 增强的内核视图 (Enhanced Kernel View)：新增列用于显示网格 (grid) 和块 (block) 大小，以及每个线程的共享内存使用量和寄存器数量。
  • 分布式训练：现已支持使用 Gloo 进行分布式训练。
  • 关联前向与反向传递中的算子：这有助于在跟踪视图 (trace view) 中将前向传递中发现的算子与反向传递关联起来，反之亦然。
  • TensorCore：该工具显示 Tensor Core (TC) 的使用情况，并为数据科学家和框架开发者提供建议。
  • NVTX：从旧版 autograd profiler 移植了对 NVTX 标记的支持。
  • 支持移动端性能分析：PyTorch Profiler 现已更好地集成 TorchScript 和移动端后端，能够为移动工作负载收集跟踪信息。

请参阅此 文档 获取详情。查看此 教程 了解如何开始使用此功能。

PyTorch Mobile

(测试版) Android NNAPI 支持进入测试阶段

去年，我们发布了 Android 神经网络 API (NNAPI) 的原型支持。NNAPI 允许 Android 应用在驱动移动设备的芯片中最强大、最高效的部分（包括 GPU 和专用 NPU）上运行计算密集型神经网络。

自原型发布以来，我们增加了更多算子覆盖，增加了对加载时灵活形状 (load-time flexible shapes) 的支持，以及在主机上运行模型进行测试的能力。请通过此 教程 尝试该功能。

此外，物体检测示例中已添加迁移学习步骤。查看此 GitHub 页面 了解更多。请在 论坛 上提供您的反馈或提问。您还可以查看 此演示文稿 以获取概览。

感谢阅读。如果您对这些更新感兴趣并希望加入 PyTorch 社区，我们鼓励您加入 讨论论坛 并在 GitHub 上提问。要获取 PyTorch 的最新消息，请关注我们的 Twitter、Medium、YouTube 和 LinkedIn。

祝好！PyTorch 团队