PyTorch 1.12 中的新库更新

由 PyTorch 团队

我们正在将多项改进引入到当前的 PyTorch 库中，同时发布 PyTorch 1.12 版本。这些更新展示了我们致力于开发跨所有领域的通用和可扩展的 API，以使我们的社区更容易在 PyTorch 上构建生态系统项目。

摘要：

• TorchVision - 添加了多权重支持 API、新的架构、模型变体和预训练权重。请参阅此处发布说明。
• TorchAudio - 引入了包括流式 API、CTC 波束搜索解码器和新的波束成形模块和方法在内的 beta 功能。请参阅此处发布说明。
• TorchText - 扩展了对可脚本化 BERT 分词器的支持，并添加了 GLUE 基准测试数据集。请参阅发布说明。
• TorchRec - 添加了 EmbeddingModule 基准测试、TwoTower 检索示例、推理和顺序嵌入、指标、改进的规划器，并展示了与生产组件的集成。请参阅发布说明。
• TorchX - 将在本地工作区开发的 PyTorch 训练器部署到五种不同的调度器上。请参阅发布说明。
• FBGemm - 为推荐系统推理工作负载添加和改进了内核，包括表批处理嵌入包、锯齿形张量操作和其他特殊情况优化。

TorchVision v0.13

多权重支持 API

TorchVision v0.13 提供了新的多权重支持 API，用于加载不同权重到现有的模型构建方法中：

from torchvision.models import *

# Old weights with accuracy 76.130%
resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)

# New weights with accuracy 80.858%
resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)

# Best available weights (currently alias for IMAGENET1K_V2)
# Note that these weights may change across versions
resnet50(weights=ResNet50_Weights.DEFAULT)

# Strings are also supported
resnet50(weights="IMAGENET1K_V2")

# No weights - random initialization
resnet50(weights=None)

新 API 捆绑了与权重相关的关键细节，如预处理转换和元数据（如标签）。以下是充分利用它的方法：

from torchvision.io import read_image
from torchvision.models import resnet50, ResNet50_Weights

img = read_image("test/assets/encode_jpeg/grace_hopper_517x606.jpg")

# Step 1: Initialize model with the best available weights
weights = ResNet50_Weights.DEFAULT
model = resnet50(weights=weights)
model.eval()

# Step 2: Initialize the inference transforms
preprocess = weights.transforms()

# Step 3: Apply inference preprocessing transforms
batch = preprocess(img).unsqueeze(0)

# Step 4: Use the model and print the predicted category
prediction = model(batch).squeeze(0).softmax(0)
class_id = prediction.argmax().item()
score = prediction[class_id].item()
category_name = weights.meta["categories"][class_id]
print(f"{category_name}: {100 * score:.1f}%")

你可以在文档中了解更多关于新 API 的信息。请使用此专门的 Github 问题来提供你的反馈。

新架构和模型变体

分类

Swin Transformer 和 EfficienetNetV2 是两种流行的分类模型，常用于下游视觉任务。本版本包括 6 个预训练权重用于它们的分类变体。以下是使用新模型的方法：

import torch
from torchvision.models import *

image = torch.rand(1, 3, 224, 224)
model = swin_t(weights="DEFAULT").eval()
prediction = model(image)

image = torch.rand(1, 3, 384, 384)
model = efficientnet_v2_s(weights="DEFAULT").eval()
prediction = model(image)

除了上述内容，我们还为现有的架构如 ShuffleNetV2、ResNeXt 和 MNASNet 提供了新的变体。以下是在 ImageNet-1K 上获得的所有新预训练模型的准确率：

我们想感谢胡叶为 TorchVision 贡献了 Swin Transformer 的实现。

(BETA) 目标检测与实例分割

我们为 RetinaNet、FasterRCNN 和 MaskRCNN 引入了 3 个新的模型变体，这些变体包括多个论文后的架构优化和改进的训练方案。所有模型都可以类似使用：

import torch
from torchvision.models.detection import *

images = [torch.rand(3, 800, 600)]
model = retinanet_resnet50_fpn_v2(weights="DEFAULT")
# model = fasterrcnn_resnet50_fpn_v2(weights="DEFAULT")
# model = maskrcnn_resnet50_fpn_v2(weights="DEFAULT")
model.eval()
prediction = model(images)

下面我们展示了新变体在 COCO val2017 上的指标。括号内表示与旧变体的改进：

我们要感谢 Ross Girshick、Piotr Dollar、Vaibhav Aggarwal、Francisco Massa 和 Hu Ye 过去对该工作的研究贡献。

新的预训练权重

SWAG 权重

ViT 和 RegNet 模型变体提供了新的预训练 SWAG（从 hashtAGs 进行弱监督）权重。其中最大的一个模型在 ImageNet-1K 上达到了惊人的 88.6%准确率。我们目前提供两种权重的版本：1）在 ImageNet-1K 上微调的端到端权重（最高准确率）和 2）冻结的树干权重，带有在 ImageNet-1K 上拟合的线性分类器（非常适合迁移学习）。下面是每个模型变体的详细准确率：

SWAG 权重在署名-非商业性使用 4.0 国际许可下发布。我们想感谢劳拉·古斯塔夫森、曼纳特·辛格和艾伦·阿德科克在使权重对 TorchVision 可用方面的工作和支持。

模型更新

多权重支持 API 的发布使我们能够刷新最受欢迎的模型并提供更精确的权重。我们在平均每个模型上提高了约 3 分。新食谱是在 ResNet50 之上学习的，其细节在之前的博客文章中已有介绍。

我们要感谢 Piotr Dollar、Mannat Singh 和 Hugo Touvron 过去对这项工作的研究和贡献。

新增增强、层和损失

本次发布带来了一大批新的基本元素，可用于构建 SOTA 模型。一些亮点包括添加 AugMix 数据增强方法、DropBlock 层、cIoU/dIoU 损失等。我们感谢 Aditya Oke、Abhijit Deo、Yassine Alouini 和 Hu Ye 为项目做出贡献，并帮助我们保持 TorchVision 的相关性和新鲜度。

文档

我们完全重写了模型文档，使其更易于浏览，并添加了各种关键信息，如支持的图像大小或预训练权重的图像预处理步骤。我们现在有一个主模型页面，包含各种可用权重的摘要表格，每个模型都有一个专门的页面。每个模型构建器也有自己的页面，提供了更多关于可用权重的详细信息，包括准确度、最小图像大小、训练食谱链接以及其他有价值的信息。为了比较，我们之前的模型文档在这里。如对新的文档提供反馈，请使用专门的 GitHub 问题。

TorchAudio v0.12

(BETA) 流式 API

StreamReader 是 TorchAudio 的新 I/O API。它由 FFmpeg†支持，并允许用户：

• 解码音频和视频格式，包括 MP4 和 AAC
• 处理输入表单，例如本地文件、网络协议、麦克风、摄像头、屏幕捕获和类似文件的对象
• 逐块迭代并解码，同时更改采样率或帧率
• 应用音频和视频过滤器，例如低通滤波器和图像缩放
• 使用 Nvidia 的基于硬件的解码器（NVDEC）解码视频

请查看文档和教程以获取使用详情：

• 媒体流 API - 第 1 部分
• 媒体流 API - 第 2 部分
• 基于 Emformer RNN-T 的在线 ASR
• 基于 Emformer RNN-T 的沉浸式设备语音识别
• 使用 NVDEC 加速视频解码

使用 StreamReader 需要 FFmpeg 库。请安装 FFmpeg。编解码器的支持范围取决于这些库的配置。TorchAudio 官方二进制文件是为与 FFmpeg 4 库兼容而编译的；如果从源代码构建 TorchAudio，则可以使用 FFmpeg 5。

(BETA) CTC 束搜索解码器

TorchAudio 集成了来自 Flashlight（GitHub）的 wav2letter CTC beam search 解码器。新增的推理时间解码器使得可以使用 TorchAudio 工具进行端到端 CTC ASR 评估。

支持自定义词典和无词典解码器，两者均兼容 KenLM n-gram 语言模型或无需使用语言模型。TorchAudio 还支持下载 LibriSpeech 数据集的 token、词典和预训练的 KenLM 文件。

使用详情请参阅文档和 ASR 推理教程。

（BETA）新的波束成形模块和方法

为了提高使用灵活性，本次发布在 torchaudio.transforms 下增加了两个新的波束成形模块：SoudenMVDR 和 RTFMVDR。与 MVDR 相比的主要区别是：

• 使用功率谱密度（PSD）和相对传递函数（RTF）矩阵作为输入，而不是时频掩码。该模块可以与直接预测语音和噪声复值短时傅里叶变换（STFT）系数的神经网络集成
• 在前向方法中添加了'reference_channel'作为输入参数，以允许用户在模型训练中选择参考通道或在推理过程中动态更改参考通道

除了这两个模块外，还在 torchaudio.functional 下增加了新的函数级波束成形方法，包括：

• psd
• mvdr_weights_souden
• mvdr_weights_rtf
• rtf_evd
• rtf_power
• apply_beamforming

使用详情请查看 torchaudio.transforms 和 torchaudio.functional 文档以及 MVDR 波束成形语音增强教程。

TorchText v0.13

粘合数据集

我们将 TorchText 中的数据集数量从 22 个增加到 30 个，通过添加 GLUE 基准测试中剩余的 8 个数据集（SST-2 已经支持）。GLUE 数据集的完整列表如下：

• CoLA（论文）：单句二元分类可接受性任务
• SST-2（论文）：单句二元分类情感任务
• MRPC（论文）：双句二元分类同义任务
• QQP：双句二元分类同义任务
• STS-B（论文）：单句到浮点回归句子相似度任务
• MNLI（论文）：句子三元分类 NLI 任务
• QNLI（论文）：句子二元分类问答和 NLI 任务
• RTE（论文）：双句子二元分类 NLI 任务
• WNLI（论文）：双句子二元分类指代和 NLI 任务

可脚本化的 BERT 分词器

TorchText 通过添加 BERT 中使用的 WordPiece 分词器扩展了对可脚本化分词器的支持。这是将输入文本分割成子词单元的常用算法之一，并在 2012 年的日语和韩语音搜索（Schuster 等人）中提出。

TorchScript 支持将 BERT 文本预处理原生嵌入 C++中，无需 Python 运行时支持。由于 TorchText 现在支持 CMAKE 构建系统，可以将 torchtext 二进制文件与应用程序代码原生链接，用户可以轻松集成 BERT 分词器以满足部署需求。

关于使用详情，请参阅相应文档。

火炬推荐 v0.2.0

嵌入模块 + DLRM 基准测试

一组基准测试，展示了 TorchRec 基础模块和研究模型的性能特征。

带有 FAISS 的 TwoTower 检索示例

我们提供了一个示例，展示了使用 TorchRec 训练一个分布式 TwoTower（即用户-物品）检索模型的过程，该模型使用分片技术。预测的物品嵌入被添加到 IVFPQ FAISS 索引中用于候选生成。检索模型和 KNN 查找被封装在一个 Pytorch 模型中，以实现高效的端到端检索。

集成

我们演示了 TorchRec 与许多在生产环境中与 PyTorch 模型一起使用的常用组件无缝配合，例如系统

• 在 Ray Clusters 上使用 Torchx Ray 调度器训练 TorchRec 模型
• 在 DLRM 上使用 NVTabular 进行预处理和数据加载
• 使用 TorchArrow 进行即时预处理的 TorchRec 模型训练，展示 RecSys 领域 UDFs

顺序嵌入示例：Bert4Rec

我们提供了一个示例，使用 TorchRec 重新实现了 BERT4REC 论文，展示了用于非池化嵌入的 EmbeddingCollection。通过 DistributedModelParallel，我们看到了相较于传统数据并行 35%的 QPS 提升。

（Beta）规划器

TorchRec 库内置了一个规划器，该规划器为给定模型选择近最优的划分方案。规划器通过评估一系列静态分析并输入整数划分器的提案来尝试识别最佳的划分方案。规划器能够自动调整适用于广泛硬件设置的方案，使用户能够无缝地从本地开发环境扩展到大规模生产硬件。请参阅此笔记本以获取更详细的教程。

(Beta) 推理

TorchRec 推理是一个支持多 GPU 推理的 C++库。TorchRec 库用于通过 torch.package（TorchScript 的替代方案）编写和打包的 Python 模型进行模型划分。torch.deploy 库用于通过启动携带打包模型的多个 Python 解释器从 C++中提供推理服务，从而绕过全局解释器锁（GIL）。提供了两个示例模型：DLRM 多 GPU（通过 TorchRec 划分）和 DLRM 单 GPU。

(Beta) RecMetrics

RecMetrics 是一个用于收集推荐模型常用工具和优化的指标库。它扩展了 torchmetrics。

• 一个集中式的指标模块，允许用户添加新的指标
• 常用指标，包括 AUC、校准、CTR、MSE/RMSE、NE 和吞吐量
• 对相关指标操作进行优化，以减少指标计算的开销
• 检查点

（原型）单进程批量+融合嵌入

之前 TorchRec 的抽象（EmbeddingBagCollection/EmbeddingCollection）在 FBGEMM 内核上，提供了诸如表批量、优化器融合和 UVM 放置等好处，但只能与 DistributedModelParallel 一起使用。我们已经将这些概念与分片解耦，并引入了 FusedEmbeddingBagCollection，它可以作为一个独立的模块使用，具有所有上述功能，也可以进行分片。

TorchX v0.2.0

TorchX 是一个作业启动器，它使得在具有许多调度器集成（包括 Kubernetes 和 Slurm）的分布式训练集群中运行 PyTorch 变得更加容易。我们很高兴发布 TorchX 0.2.0 版本，其中包含了许多改进。TorchX 目前正在本地和云环境中投入生产使用。

查看快速入门指南以开始启动本地和远程作业。

工作空间

TorchX 现在支持工作空间，使用户能够轻松地使用本地工作空间启动训练作业。TorchX 可以自动在基础镜像上构建包含您本地训练代码的补丁，以最小化迭代时间和训练时间。

.torchxconfig

在 .torchxconfig 中指定选项可以让你每次启动作业时不必输入冗长的 CLI 命令。你还可以定义项目级别的通用配置，并在用户目录中放置配置文件以进行用户级别的覆盖。

扩展调度器支持

除了现有的集成之外，TorchX 现在还支持 AWS Batch 和 Ray（实验性）调度器。

所有调度器的分布式训练

现在 TorchX 的 dist.ddp 组件可以在所有调度器上工作，无需任何配置。使用 torchelastic 的内置 dist.ddp 组件时，分布式训练工作进程将自动发现彼此。

超参数优化

TorchX 与 Ax 集成，让您可以通过在远程集群上启动搜索试验来扩展超参数优化（HPO）。

文件和设备挂载

TorchX 现在支持远程文件系统挂载和自定义设备。这使得您的 PyTorch 作业可以高效地访问云存储，如 NFS 或 Lustre。设备挂载功能使得可以使用网络加速器，如 Infiniband 以及自定义推理/训练加速器。

FBGemm v0.2.0

FBGEMM 库包含旨在提高 PyTorch 工作负载性能的优化内核。在过去几个月中，我们添加了许多新功能和优化，我们非常高兴向大家报告。

推理表批量嵌入（TBE）

表批量嵌入包（TBE）算子是推荐系统推理在 GPU 上进行嵌入查找的重要基础操作。为了提高性能和灵活性，我们添加了以下增强功能：

移除了对齐限制

• 嵌入维度 * 数据类型大小之前必须是 4B 的倍数，现在则是 1B。

统一虚拟内存（UVM）缓存内核优化

• UVM 缓存内核现在与使用 UVM 缓存的表数量成线性扩展。之前，它的开销与所有使用 UVM 缓存的表相似
• UVM 缓存内核的开销比之前小得多

推理 FP8 表批量嵌入（TBE）

表批处理嵌入袋（TBE）之前支持 FP32、FP16、INT8、INT4 和 INT2 嵌入权重类型。虽然这些权重类型在许多模型中表现良好，但我们集成了 FP8 权重类型（在 GPU 和 CPU 操作中），以便在我们的模型中对 FP8 进行数值和性能评估。与 INT8 相比，FP8 不需要额外的偏差和尺度存储和计算。此外，下一代 H100 GPU 在 Tensor Core（主要是矩阵乘法操作）上支持 FP8。

锯齿形张量内核

我们添加了优化内核以加速 TorchRec 的 JaggedTensor。JaggedTensor 的目的是处理输入数据的一个维度是“锯齿形”的情况，这意味着给定维度中连续的行在长度上可能不同，这在推荐系统中的稀疏特征输入中很常见。内部表示如下：

我们添加了将锯齿形张量从稀疏格式转换为密集格式以及反向转换的算子，以及与锯齿形张量进行矩阵乘法和逐元素操作的算子。

优化后的 permute102-baddbmm-permute102

在批大小不是模型批大小的情况下，融合各种矩阵乘法比较困难，切换批维度是一个快速解决方案。我们创建了 permute102_baddbmm_permute102 操作，该操作切换第一维和第二维，执行批矩阵乘法，然后切换回来。目前我们只支持 FP16 数据类型的正向传播，并将支持 FP32 类型和反向传播。

优化后的 dim 0 索引选择的 index_select

index_select 通常用作稀疏操作的一部分。虽然 PyTorch 支持任意维度索引选择的通用 index_select，但其对 dim 0 索引选择这类特殊情况的性能不佳。因此，我们实现了一个针对 dim 0 的专用 index_select。在某些情况下，我们观察到 FBGEMM 的 index_select 比 PyTorch（使用均匀索引分布）的 index_select 性能提升了 1.4 倍。

更多关于有影响力实例的实现信息可以在我们的 GitHub 页面和教程中找到。

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

喝彩！

PyTorch 团队