除了 PyTorch 1.12 版本发布外，我们还对现有的 PyTorch 库进行了一系列改进。这些更新展示了我们致力于在所有领域开发通用且可扩展的 API，旨在让我们的社区能够更轻松地在 PyTorch 上构建生态系统项目。
总结
- TorchVision – 增加了多权重支持 API、新架构、模型变体和预训练权重。点击此处查看发布说明。
- TorchAudio – 引入了测试阶段（Beta）功能，包括流式传输 API、CTC 波束搜索解码器，以及新的波束形成模块和方法。点击此处查看发布说明。
- TorchText – 扩展了对可脚本化 BERT 分词器的支持，并添加了用于 GLUE 基准测试的数据集。点击此处查看发布说明。
- TorchRec – 增加了 EmbeddingModule 基准测试，提供了 TwoTower 检索、推理和序列嵌入、指标的示例，改进了规划器，并展示了与生产环境组件的集成。点击此处查看发布说明。
- TorchX – 可将本地工作区开发的 PyTorch 训练任务启动到五种不同类型的调度器上。点击此处查看发布说明。
- FBGemm – 为推荐系统推理工作负载添加并改进了内核，包括表批处理嵌入包（table batched embedding bag）、锯齿状张量（jagged tensor）运算以及其他特殊情况优化。
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 将预处理变换（preprocessing transforms）和标签等元数据与权重打包在一起。以下是如何充分利用该功能的方法。
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 Issue。
新架构和模型变体
分类
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 上获得的准确率。
| 模型 | Acc@1 | Acc@5 |
|---|---|---|
| swin_t | 81.474 | 95.776 |
| swin_s | 83.196 | 96.36 |
| swin_b | 83.582 | 96.64 |
| efficientnet_v2_s | 84.228 | 96.878 |
| efficientnet_v2_m | 85.112 | 97.156 |
| efficientnet_v2_l | 85.808 | 97.788 |
| resnext101_64x4d | 83.246 | 96.454 |
| resnext101_64x4d (量化) | 82.898 | 96.326 |
| shufflenet_v2_x1_5 | 72.996 | 91.086 |
| shufflenet_v2_x1_5 (量化) | 72.052 | 0.700 |
| shufflenet_v2_x2_0 | 76.230 | 93.006 |
| shufflenet_v2_x2_0 (量化) | 75.354 | 92.488 |
| mnasnet0_75 | 71.180 | 90.496 |
| mnas1_3 | 76.506 | 93.522 |
我们要感谢 Hu Ye 为 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 数据集上的指标。括号内表示相较于旧变体的提升幅度。
| 模型 | Box mAP | Mask mAP |
|---|---|---|
| retinanet_resnet50_fpn_v2 | 41.5 (+5.1) | – |
| fasterrcnn_resnet50_fpn_v2 | 46.7 (+9.7) | – |
| maskrcnn_resnet50_fpn_v2 | 47.4 (+9.5) | 41.8 (+7.2) |
我们要感谢 Ross Girshick、Piotr Dollar、Vaibhav Aggarwal、Francisco Massa 和 Hu Ye 的既往研究及对该工作的贡献。
新的预训练权重
SWAG 权重
ViT 和 RegNet 模型变体提供了新的预训练 SWAG（Supervised Weakly from hashtAGs）权重。其中最大的模型在 ImageNet-1K 上达到了惊人的 88.6% 的准确率。我们目前提供两个版本的权重：1) 在 ImageNet-1K 上进行端到端微调的权重（准确率最高）；2) 在 ImageNet-1K 上通过线性分类器拟合的冻结主干权重（非常适合迁移学习）。以下是各模型变体的详细准确率。
| 模型权重 | Acc@1 | Acc@5 |
|---|---|---|
| RegNet_Y_16GF_Weights.IMAGENET1K_SWAG_E2E_V1 | 86.012 | 98.054 |
| RegNet_Y_16GF_Weights.IMAGENET1K_SWAG_LINEAR_V1 | 83.976 | 97.244 |
| RegNet_Y_32GF_Weights.IMAGENET1K_SWAG_E2E_V1 | 86.838 | 98.362 |
| RegNet_Y_32GF_Weights.IMAGENET1K_SWAG_LINEAR_V1 | 84.622 | 97.48 |
| RegNet_Y_128GF_Weights.IMAGENET1K_SWAG_E2E_V1 | 88.228 | 98.682 |
| RegNet_Y_128GF_Weights.IMAGENET1K_SWAG_LINEAR_V1 | 86.068 | 97.844 |
| ViT_B_16_Weights.IMAGENET1K_SWAG_E2E_V1 | 85.304 | 97.65 |
| ViT_B_16_Weights.IMAGENET1K_SWAG_LINEAR_V1 | 81.886 | 96.18 |
| ViT_L_16_Weights.IMAGENET1K_SWAG_E2E_V1 | 88.064 | 98.512 |
| ViT_L_16_Weights.IMAGENET1K_SWAG_LINEAR_V1 | 85.146 | 97.422 |
| ViT_H_14_Weights.IMAGENET1K_SWAG_E2E_V1 | 88.552 | 98.694 |
| ViT_H_14_Weights.IMAGENET1K_SWAG_LINEAR_V1 | 85.708 | 97.73 |
SWAG 权重根据 署名-非商业性使用 4.0 国际 (CC BY-NC 4.0) 许可发布。我们要感谢 Laura Gustafson、Mannat Singh 和 Aaron Adcock 的工作与支持，使这些权重能够在 TorchVision 中使用。
模型更新
多权重支持 API 的发布使我们能够更新最流行的模型，并提供更准确的权重。每个模型的准确率平均提升了约 3 个百分点。所使用的新配方是基于 ResNet50 学习得来的，其细节在之前的博客文章中已有介绍。
| 模型 | 旧权重 | 新权重 |
|---|---|---|
| efficientnet_b1 | 78.642 | 79.838 |
| mobilenet_v2 | 71.878 | 72.154 |
| mobilenet_v3_large | 74.042 | 75.274 |
| regnet_y_400mf | 74.046 | 75.804 |
| regnet_y_800mf | 76.42 | 78.828 |
| regnet_y_1_6gf | 77.95 | 80.876 |
| regnet_y_3_2gf | 78.948 | 81.982 |
| regnet_y_8gf | 80.032 | 82.828 |
| regnet_y_16gf | 80.424 | 82.886 |
| regnet_y_32gf | 80.878 | 83.368 |
| regnet_x_400mf | 72.834 | 74.864 |
| regnet_x_800mf | 75.212 | 77.522 |
| regnet_x_1_6gf | 77.04 | 79.668 |
| regnet_x_3_2gf | 78.364 | 81.196 |
| regnet_x_8gf | 79.344 | 81.682 |
| regnet_x_16gf | 80.058 | 82.716 |
| regnet_x_32gf | 80.622 | 83.014 |
| resnet50 | 76.13 | 80.858 |
| resnet50 (量化) | 75.92 | 80.282 |
| resnet101 | 77.374 | 81.886 |
| resnet152 | 78.312 | 82.284 |
| resnext50_32x4d | 77.618 | 81.198 |
| resnext101_32x8d | 79.312 | 82.834 |
| resnext101_32x8d (量化) | 78.986 | 82.574 |
| wide_resnet50_2 | 78.468 | 81.602 |
| wide_resnet101_2 | 78.848 | 82.51 |
我们要感谢 Piotr Dollar、Mannat Singh 和 Hugo Touvron 的既往研究及对该工作的贡献。
新的增强、层和损失函数
此版本带来了一系列可用于生产 SOTA 模型的新原语。亮点包括添加了 AugMix 数据增强方法、DropBlock 层、cIoU/dIoU 损失函数以及更多内容。感谢 Aditya Oke、Abhijit Deo、Yassine Alouini 和 Hu Ye 对项目的贡献，并帮助我们保持 TorchVision 的相关性和前沿性。
文档
我们彻底重构了模型文档，使其更易于浏览，并添加了各种关键信息，例如支持的图像尺寸或预训练权重的图像预处理步骤。现在我们拥有一个主模型页面，其中包含各种可用权重的摘要表，每个模型都有一个专用页面。每个模型构建器也在其自己的页面中进行了说明，提供了关于可用权重的更多详情，包括准确率、最小图像尺寸、训练方案链接和其他有价值的信息。作为对比，我们之前的模型文档位于此处。如需针对新文档提供反馈，请使用专门的 Github Issue。
TorchAudio v0.12
(BETA) 流式传输 API
StreamReader 是 TorchAudio 的全新 I/O API。它基于 FFmpeg† 构建，允许用户：
- 解码音频和视频格式，包括 MP4 和 AAC
- 处理输入形式，例如本地文件、网络协议、麦克风、网络摄像头、屏幕录制和类似文件的对象
- 逐块迭代并解码，同时更改采样率或帧率
- 应用音频和视频过滤器，例如低通滤波器和图像缩放
- 利用 Nvidia 的硬件解码器 (NVDEC) 进行视频解码
有关使用详情，请查阅文档和教程
- 媒体流 API – 第一部分
- 媒体流 API – 第二部分
- 使用 Emformer RNN-T 进行在线 ASR
- 使用 Emformer RNN-T 进行设备端 ASR
- 使用 NVDEC 进行加速视频解码
† 要使用 StreamReader，需要 FFmpeg 库。请安装 FFmpeg。编解码器的覆盖范围取决于这些库的配置方式。TorchAudio 官方二进制文件编译为与 FFmpeg 4 库配合使用；如果从源代码构建 TorchAudio，则可以使用 FFmpeg 5。
(BETA) CTC 波束搜索解码器
TorchAudio 集成了来自 Flashlight (GitHub) 的 wav2letter CTC 波束搜索解码器。该推理阶段解码器的加入，使得使用 TorchAudio 工具运行端到端 CTC ASR 评估成为可能。
支持可定制的词典和无词典解码器，两者均与 KenLM n-gram 语言模型兼容（也可不使用语言模型）。TorchAudio 此外还支持为 LibriSpeech 数据集下载 token、词典和预训练的 KenLM 文件。
(BETA) 新的波束形成模块和方法
为了提高使用的灵活性，此版本在 torchaudio.transforms 下添加了两个新的波束形成模块：SoudenMVDR 和 RTFMVDR。它们与 MVDR 的主要区别在于：
- 使用功率谱密度 (PSD) 和相对传递函数 (RTF) 矩阵作为输入，而不是时频掩码。该模块可以集成到直接预测语音和噪声的复数 STFT 系数的神经网络中。
- 在 forward 方法中将 ‘reference_channel’ 添加为输入参数，允许用户在模型训练中选择参考通道，或在推理中动态更改参考通道。
除了这两个模块外，在 torchaudio.functional 下还添加了新的函数级波束形成方法。这些包括：
有关使用详情，请查看 torchaudio.transforms 和 torchaudio.functional 的文档，以及 使用 MVDR 波束形成进行语音增强的教程。
TorchText v0.13
Glue 数据集
我们通过添加 GLUE 基准测试中剩余的 8 个数据集，将 TorchText 中的数据集数量从 22 个增加到了 30 个（SST-2 已被支持）。GLUE 数据集的完整列表如下：
- CoLA (论文)：单句二分类可接受性任务
- SST-2 (论文)：单句二分类情感任务
- MRPC (论文)：双句二分类释义任务
- QQP：双句二分类释义任务
- STS-B (论文)：单句到浮点数回归句子相似度任务
- MNLI (论文)：句子三分类 NLI 任务
- QNLI (论文)：句子二分类 QA 和 NLI 任务
- RTE (论文)：双句二分类 NLI 任务
- WNLI (论文)：双句二分类指代消解和 NLI 任务
可脚本化 BERT 分词器
TorchText 通过添加 BERT 中使用的 WordPiece 分词器，扩展了对可脚本化分词器的支持。这是将输入文本拆分为子词单元最常用的算法之一，最初在 日语和韩语语音搜索 (Schuster et al., 2012) 中引入。
TorchScript 的支持将允许用户在 C++ 中原生嵌入 BERT 文本预处理，而无需 Python 运行时的支持。由于 TorchText 现在支持 CMAKE 构建系统，可将 torchtext 二进制文件与应用程序代码原生链接，用户可以轻松集成 BERT 分词器以满足部署需求。
有关使用详情，请参阅相应的 文档。
TorchRec v0.2.0
EmbeddingModule + DLRM 基准测试
我们提供了一套 基准测试，展示了 TorchRec 基础模块和基于 TorchRec 构建的研究模型的性能特征。
TwoTower 检索示例（使用 FAISS）
我们提供了一个 示例，演示了如何训练使用 TorchRec 进行分片的分布式 TwoTower（即用户-物品）检索模型。投影后的物品嵌入被添加到 IVFPQ FAISS 索引中以进行候选生成。检索模型和 KNN 查找被捆绑在同一个 Pytorch 模型中，以实现高效的端到端检索。
集成
我们展示了 TorchRec 可以与生产系统中常与 PyTorch 模型配合使用的许多组件实现“开箱即用”，例如：
- 训练：使用 Torchx Ray 调度器在 Ray 集群上训练 TorchRec 模型
- 预处理：在 DLRM 上使用 NVTabular 进行预处理和数据加载
- 训练：使用 TorchArrow 进行即时预处理来训练 TorchRec 模型，展示了 RecSys 领域的 UDF
序列嵌入示例：Bert4Rec
我们提供了一个使用 TorchRec 的 示例，它重新实现了 BERT4REC 论文，展示了用于非池化嵌入的 EmbeddingCollection。通过使用 DistributedModelParallel，我们看到了相较于传统数据并行 35% 的 QPS 提升。
(Beta) 规划器 (Planner)
TorchRec 库包含一个内置的 规划器，它为给定模型选择近乎最优的分片方案。规划器尝试通过评估一系列经过静态分析并输入到整数分区器中的方案来确定最佳分片计划。规划器能够针对各种硬件设置自动调整计划，使用户能够从本地开发环境到大规模生产硬件无缝扩展性能。有关更详细的教程，请参见此 笔记本。
(Beta) 推理
TorchRec Inference 是一个支持多 GPU 推理的 C++ 库。TorchRec 库用于对通过 torch.package（TorchScript 的替代品）编写和打包的 Python 模型进行分片。torch.deploy 库通过启动承载打包模型的多个 Python 解释器来从 C++ 提供推理服务，从而规避了 GIL（全局解释器锁）。提供了两个模型作为示例：DLRM 多 GPU（通过 TorchRec 分片）和 DLRM 单 GPU。
(Beta) RecMetrics
RecMetrics 是一个 指标 库，它收集了推荐模型的常见工具和优化方法。它扩展了 torchmetrics。
- 一个允许用户添加新指标的集中式指标模块
- 常用指标，包括 AUC、Calibration、CTR、MSE/RMSE、NE 和吞吐量
- 针对指标相关操作的优化，以减少指标计算的开销
- 检查点 (Checkpointing)
(Prototype) 单进程批处理 + 融合嵌入
此前，TorchRec 在 FBGEMM 内核之上的抽象（EmbeddingBagCollection/EmbeddingCollection）——提供表批处理、优化器融合和 UVM 放置等优势——只能与 DistributedModelParallel 结合使用。我们已将这些概念与分片解耦，并引入了 FusedEmbeddingBagCollection，它既可以作为独立模块使用（具备上述所有功能），也可以进行分片。
TorchX v0.2.0
TorchX 是一个作业启动器，旨在让在分布式训练集群（包括 Kubernetes 和 Slurm）上运行 PyTorch 变得更加容易。我们很高兴发布 TorchX 0.2.0，它带来了多项改进。TorchX 目前正被用于内部部署和云环境的生产中。
查看 快速入门 以开始启动本地和远程作业。
工作区 (Workspaces)
TorchX 现已支持工作区，允许用户轻松使用其本地工作区启动训练任务。TorchX 可以自动在基础镜像之上构建一个包含本地训练代码的补丁，以最小化迭代时间和训练启动时间。
.torchxconfig
在 .torchxconfig 中指定选项可省去您每次启动任务时输入冗长 CLI 命令的麻烦。您还可以定义项目级通用配置，并在主目录中放入一个配置文件以进行用户级覆盖。
扩展的调度器支持
除现有的集成外，TorchX 现在还支持 AWS Batch 和 Ray (实验性) 调度器。
所有调度器上的分布式训练
TorchX dist.ddp 组件现在可以在所有调度器上运行，无需任何配置。当使用 内置的 dist.ddp 组件 通过 torchelastic 时，分布式训练工作节点将自动互相发现。
超参数优化
TorchX 与 Ax 集成，让您可以通过将搜索试验启动到远程集群来扩展超参数优化 (HPO)。
文件和设备挂载
TorchX 现在支持 远程文件系统挂载和自定义设备。这使您的 PyTorch 作业能够高效访问 NFS 或 Lustre 等云存储。设备挂载则支持使用 Infiniband 等网络加速器以及自定义推理/训练加速器。
FBGemm v0.2.0
FBGEMM 库包含旨在提高 PyTorch 工作负载性能的优化内核。在过去几个月中，我们添加了许多新功能和优化，很高兴在此向大家汇报。
推理表批处理嵌入 (TBE)
表批处理嵌入包 (TBE) 算子是 GPU 上推荐系统推理嵌入查找的重要基础操作。我们针对性能和灵活性进行了以下增强：
移除了对齐限制
- 以前要求“嵌入维度 * 数据类型大小”必须是 4 字节的倍数，现在变为 1 字节。
统一虚拟内存 (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
当 batch size 不是模型 batch size 时，融合各种矩阵乘法很困难；切换批处理维度是一个快速解决方案。我们创建了 permute102_baddbmm_permute102 运算，它交换第一和第二个维度，执行批处理矩阵乘法，然后切回。目前我们仅支持 FP16 数据类型的前向传播，未来将支持 FP32 类型和反向传播。
针对维度 0 索引选择优化 index_select
index_select 通常用作稀疏操作的一部分。虽然 PyTorch 支持用于任意维度索引选择的通用 index_select，但其在维度 0 索引选择等特殊情况下的性能并不理想。因此，我们实现了针对维度 0 的专业化 index_select。在某些情况下，我们观察到 FBGEMM 的 index_select 相比 PyTorch 的实现（使用均匀索引分布）性能提升了 1.4 倍。
有关有影响力实例的实现的更多信息，请访问我们的 GitHub 页面和教程。
感谢阅读，如果您对这些更新感兴趣并希望加入 PyTorch 社区，我们鼓励您加入 讨论论坛 并 提交 GitHub issue。要获取 PyTorch 的最新消息，请在 Twitter、Medium、YouTube 和 LinkedIn 上关注我们。
干杯！
PyTorch 团队