PyTorch 1.12：TorchArrow、模块的功能性 API 和 nvFuser 现已可用

我们很高兴地宣布发布 PyTorch 1.12（发布说明）！本次发布包含 3124 多次提交，由 433 位贡献者完成。除了 1.12 之外，我们还发布了 AWS S3 集成、CPU 上 Channels Last 的 PyTorch Vision 模型、使用 Bfloat16 和 FSDP API 增强 Intel® Xeon® 可扩展处理器上的 PyTorch 的 Beta 版本。我们衷心感谢我们敬业的社区所做的贡献。

总结

功能性 API，用于使用给定参数集功能性地应用模块计算
PyTorch 中的 Complex32 和复杂卷积
TorchData 的 DataPipes 完全向后兼容 DataLoader
functorch 改进了 API 覆盖范围
nvFuser，一个用于 PyTorch 的深度学习编译器
Ampere 及更高版本 CUDA 硬件上 float32 矩阵乘法精度的更改
TorchArrow，一个用于批处理数据的机器学习预处理新测试版库

前端 API

介绍 TorchArrow

我们已经准备好了一个新的 Beta 版本供您尝试和使用：TorchArrow。这是一个用于批处理数据的机器学习预处理库。它具有高性能的 Pandas 风格、易于使用的 API，以加快您的预处理工作流程和开发。

目前，它提供了具有以下功能的 Python DataFrame 接口

高性能 CPU 后端，使用 Velox 的矢量化和可扩展的用户定义函数 (UDF)
与 PyTorch 或其他模型创作无缝衔接，例如张量整理以及轻松插入 PyTorch DataLoader 和 DataPipes
通过 Arrow 内存列式格式实现外部读取器的零拷贝

有关更多详细信息，请参阅我们的 10 分钟教程、安装说明、API 文档，以及 TorchRec 中数据预处理的原型。

(Beta) 模块的功能性 API

PyTorch 1.12 引入了一个新的 Beta 功能，用于使用给定参数集功能性地应用模块计算。有时，传统的 PyTorch 模块使用模式（在内部维护一组静态参数）过于严格。在实现元学习算法时经常出现这种情况，其中可能需要在优化器步骤中维护多组参数。

新的 torch.nn.utils.stateless.functional_call() API 允许

模块计算对所使用的参数集具有完全的灵活性
无需以功能性方式重新实现您的模块
模块中存在的任何参数或缓冲区都可以与外部定义的值交换以供调用使用。引用参数/缓冲区的命名遵循模块 state_dict() 中的完全限定形式

例子

import torch
from torch import nn
from torch.nn.utils.stateless import functional_call

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(3, 3)
        self.bn = nn.BatchNorm1d(3)
        self.fc2 = nn.Linear(3, 3)

    def forward(self, x):
        return self.fc2(self.bn(self.fc1(x)))

m = MyModule()

# Define parameter / buffer values to use during module computation.
my_weight = torch.randn(3, 3, requires_grad=True)
my_bias = torch.tensor([1., 2., 3.], requires_grad=True)
params_and_buffers = {
    'fc1.weight': my_weight,
    'fc1.bias': my_bias,
    # Custom buffer values can be used too.
    'bn.running_mean': torch.randn(3),
}

# Apply module computation to the input with the specified parameters / buffers.
inp = torch.randn(5, 3)
output = functional_call(m, params_and_buffers, inp)

(Beta) PyTorch 中的 Complex32 和复杂卷积

PyTorch 今天原生支持复数、复杂自动梯度、复杂模块和许多复杂操作，包括线性代数和快速傅里叶变换 (FFT) 运算符。许多库，包括 torchaudio 和 ESPNet，已经使用 PyTorch 中的复数，PyTorch 1.12 通过复杂卷积和实验性 complex32（“复杂半精度”）数据类型进一步扩展了复杂功能，该数据类型支持半精度 FFT 操作。由于 CUDA 11.3 包中的错误，如果您使用复数，我们建议使用 wheels 中的 CUDA 11.6 包。

(Beta) 正向模式自动微分

正向模式 AD 允许在正向传播中及时计算方向导数（或等效地，雅可比向量积）。PyTorch 1.12 显著提高了正向模式 AD 的运算符覆盖范围。有关更多信息，请参阅我们的教程。

TorchData

BC DataLoader + DataPipe

TorchData 中的“DataPipe”在多进程和分布式环境中，在洗牌确定性和动态分片方面完全向后兼容现有的“DataLoader”。

(Beta) AWS S3 集成

基于 AWSSDK 的 DataPipes 已集成到 TorchData 中。它提供了由原生 AWSSDK 支持的以下功能

根据前缀从每个 S3 存储桶检索 URL 列表
- 支持超时以防止无限期挂起
- 支持指定 S3 存储桶区域
从 S3 URL 加载数据
- 支持缓冲和多部分下载
- 支持指定 S3 存储桶区域

AWS 原生 DataPipes 仍处于 Beta 阶段。我们将继续对其进行调整以提高其性能。

(原型) DataLoader2

DataLoader2 以原型模式提供。我们正在引入 DataPipes、DataLoading API 和后端（即 ReadingServices）之间交互的新方式。该功能在 API 方面是稳定的，但在功能上尚未完善。我们欢迎早期采用者和反馈，以及潜在的贡献者。

有关更多详细信息，请查看此链接。

functorch

受 Google JAX 启发，functorch 是一个提供可组合 vmap（矢量化）和自动微分变换的库。它支持在 PyTorch 中难以表达的高级自动微分用例。这些示例包括

我们很高兴地宣布 functorch 0.2.0，它包含多项改进和新的实验性功能。

显著改进的覆盖率

我们显著改进了 functorch.jvp （我们的正向模式自动微分 API）和依赖于它的其他 API（functorch.{jacfwd, hessian}）的覆盖率。

(原型) functorch.experimental.functionalize

给定函数 f，functionalize(f) 返回一个没有突变的新函数（有条件）。这对于构建没有就地操作的 PyTorch 函数跟踪很有用。例如，您可以使用 make_fx(functionalize(f)) 来构建 PyTorch 函数的无突变跟踪。要了解更多信息，请参阅文档。

有关更多详细信息，请参阅我们的安装说明、文档、教程和发布说明。

性能改进

介绍 nvFuser，一个用于 PyTorch 的深度学习编译器

在 PyTorch 1.12 中，Torchscript 正在将其默认融合器（用于 Volta 及更高版本的 CUDA 加速器）更新为 nvFuser，它支持更广泛的操作范围，并且比 NNC（以前用于 CUDA 设备的融合器）更快。即将发布的一篇博客文章将详细介绍 nvFuser，并展示它如何加速各种网络的训练。

有关使用和调试的更多详细信息，请参阅 nvFuser 文档。

Ampere 及更高版本 CUDA 硬件上 float32 矩阵乘法精度的更改

PyTorch 支持各种“混合精度”技术，例如 torch.amp（自动混合精度）模块，以及在 Ampere 及更高版本 CUDA 硬件上使用 TensorFloat32 数据类型执行 float32 矩阵乘法以加快内部计算。在 PyTorch 1.12 中，我们正在更改 float32 矩阵乘法的默认行为，使其始终使用完整的 IEEE fp32 精度，这比使用 TensorFloat32 数据类型进行内部计算更精确但更慢。对于 TensorFloat32 与 float32 吞吐量比例特别高的设备（例如 A100），这种默认更改可能会导致显着的速度下降。

如果您一直在使用 TensorFloat32 矩阵乘法，则可以通过设置 torch.backends.cuda.matmul.allow_tf32 = True 继续这样做。

这自 PyTorch 1.7 起受支持。从 PyTorch 1.12 开始，也可以使用新的 matmul 精度 API：torch.set_float32_matmul_precision(“highest”|”high”|”medium”)

重申一下，PyTorch 的新默认设置是所有设备类型的“最高”精度。我们认为这为矩阵乘法提供了更好的跨设备类型一致性。有关新精度 API 的文档可以在此处找到。“高”或“中”精度类型的设置将使 Ampere 及更高版本的 CUDA 设备上的 TensorFloat32 成为可能。如果您正在更新到 PyTorch 1.12，为了保持 Ampere 设备上矩阵乘法的当前行为和更快性能，请将精度设置为“高”。

使用混合精度技术对于高效训练许多现代深度学习网络至关重要，如果您已经在使用 torch.amp，则此更改不太可能影响您。如果您不熟悉混合精度训练，请参阅我们即将发布的“每个用户都应该了解的 PyTorch 混合精度训练”博客文章。

(Beta) 使用 Channels Last 在 CPU 上加速 PyTorch 视觉模型

在运行视觉模型时，内存格式对性能有显著影响，通常 Channels Last 从性能角度来看更具优势，因为数据局部性更好。1.12 包含了内存格式的基本概念，并展示了在 Intel® Xeon® 可扩展处理器上，在流行的 PyTorch 视觉模型上使用 Channels Last 的性能优势。

在 CPU 上为 CV 领域常用运算符启用 Channels Last 内存格式支持，适用于推理和训练
为 ATen 中的 Channels Last 内核提供原生级别优化，适用于 AVX2 和 AVX512
在 Intel® Xeon® Ice Lake（或更新版本）CPU 上，TorchVision 模型的推理性能比 Channels First 提高了 1.3 倍到 1.8 倍

(Beta) 使用 Bfloat16 增强 Intel® Xeon® 可扩展处理器上的 PyTorch

降低精度数值格式，如 bfloat16，提高了 PyTorch 在多个深度学习训练工作负载上的性能。PyTorch 1.12 包含了 bfloat16 的最新软件增强功能，适用于更广泛的用户场景，并展示了更高的性能提升。主要改进包括

使用 Intel® Xeon® Cooper Lake CPU 中引入的新 bfloat16 本机指令 VDPBF16PS，硬件计算吞吐量是 float32 的 2 倍
内存占用是 float32 的 1/2，对于内存带宽密集型运算符速度更快
在 Intel® Xeon® Cooper Lake（或更新版本）CPU 上，TorchVision 模型的推理性能比 float32 提高了 1.4 倍到 2.2 倍

(原型) 介绍在 Mac 上加速 PyTorch 训练

随着 PyTorch 1.12 的发布，开发人员和研究人员现在可以利用 Apple 芯片 GPU 来显著加快模型训练。这使得可以在本地 Mac 上执行机器学习工作流程，例如原型设计和微调。加速 GPU 训练是使用 Apple 的 Metal Performance Shaders (MPS) 作为后端实现的。其好处包括加速 GPU 训练带来的性能提升，以及在本地训练更大网络或批量大小的能力。在此处了解更多信息：这里。

加速 GPU 训练和评估速度比仅 CPU 快（倍数）

除了新的 MPS 设备支持之外，前几个版本中已提供的核心和领域库的 M1 二进制文件现在已成为官方原型功能。这些二进制文件可用于在 Apple Silicon 上原生运行 PyTorch。

(原型) BetterTransformer：Transformer 编码器推理的快速路径执行

PyTorch 现在支持 TransformerEncoder、TransformerEncoderLayer 和 MultiHeadAttention (MHA) 等多个 Transformer 编码器模块的 CPU 和 GPU 快速路径实现（“BetterTransformer”）。BetterTransformer 快速路径架构始终更快——对于许多常见的执行场景，根据模型和输入特性，速度提高 2 倍。新的 BetterTransformer 启用模块与 PyTorch Transformer API 的先前版本 API 兼容，如果它们满足快速路径执行要求，将加速现有模型，并读取使用 PyTorch 先前版本训练的模型。PyTorch 1.12 包括

Torchtext 预训练的 RoBERTa 和 XLM-R 模型的 BetterTransformer 集成
Torchtext 基于 PyTorch Transformer API 构建
通过融合内核（将多个运算符组合成一个内核）减少执行开销，从而提高性能的快速路径执行
通过在自然语言处理中处理填充标记时利用数据稀疏性来获得额外加速的选项（通过在创建 TransformerEncoder 时设置 enable_nested_tensor=True）
帮助用户了解为何未发生快速路径执行的诊断

分布式

(Beta) 完全分片数据并行 (FSDP) API

FSDP API 通过在数据并行工作器之间分片模型的参数、梯度和优化器状态，同时保持数据并行的简单性，帮助轻松扩展大型模型训练。原型版本在 PyTorch 1.11 中发布，包含最少的功能集，帮助了对多达 1T 参数的模型进行缩放测试。

在此 Beta 版本中，FSDP API 添加了以下功能以支持各种生产工作负载。此 Beta 版本中新添加功能的主要亮点包括

通用分片策略 API – 用户只需更改一行代码即可轻松切换分片策略，从而比较和使用 DDP（仅数据分片）、FSDP（完整模型和数据分片）或 Zero2（仅优化器和梯度分片）来优化其特定训练需求的内存和性能
精细混合精度策略 – 用户可以通过混合精度策略为模型参数、梯度通信和缓冲区指定半精度和全精度数据类型（bfloat16、fp16 或 fp32）的混合。模型会自动以 fp32 格式保存，以实现最大可移植性
Transformer 自动包装策略 – 通过注册模型层类实现基于 Transformer 的模型的最佳包装，从而加速训练性能
使用 device_id 初始化实现更快的模型初始化 – 初始化以流式方式执行，以避免 OOM 问题并优化初始化性能，优于 CPU 初始化
Rank0 流式传输用于大型模型的完整模型保存 – 完全分片模型可以通过所有 GPU 将其分片流式传输到 rank 0 GPU 来保存，并且模型在 rank 0 CPU 上以完整状态构建以进行保存

有关更多详细信息和示例代码，请查看文档和教程。

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

干杯！

PyTorch 团队

PyTorch 1.12：TorchArrow、模块的函数式 API 和 nvFuser 现已可用