博客

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

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

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

PyTorch 1.10 的更新重点在于提升训练性能及开发人员的易用性。完整的发布说明可点击此处查看。亮点包括：

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

随 1.10 版本一同发布的还有 PyTorch 库的重大更新，您可以在这篇博文中了解详情。

前端 API

(稳定版) 使用 FX 进行 Python 代码转换

FX 提供了一个 Python 风格的平台，用于转换和降低（lowering）PyTorch 程序。它是一个为转换编写者提供的工具包，旨在促进函数和 nn.Module 实例的 Python 到 Python 转换。该工具包的目标是支持 Python 语言语义的一个子集——而非整个 Python 语言——以便于轻松实现转换。随 1.10 版本发布，FX 正式进入稳定版。

您可以通过官方文档以及使用 torch.fx 实现程序转换的 GitHub 示例进一步了解 FX。

(稳定版) torch.special

torch.special 模块现已进入稳定版，类似于 SciPy 的 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 还降低了抖动（jitter）；由于并行工作负载必须等待最慢的 worker，因此减少抖动可以提高整体并行效率。

此次集成允许在被 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 包含以下关键功能：

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

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

PyTorch Mobile

(测试版) Android NNAPI 支持

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

自原型发布以来，我们增加了更多算子覆盖，支持了加载时的动态形状，以及在主机上运行模型进行测试的功能。请尝试使用此教程来体验该功能。

此外，我们还在目标检测示例中添加了迁移学习步骤。请查阅此 GitHub 页面了解更多。欢迎在论坛提供反馈或提问。您还可以查看此演示视频以获取概述。

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

祝好！PyTorch 团队