PyTorch 1.13 中的新库更新

摘要

随着 PyTorch 1.13 版本的发布，我们对现有的 PyTorch 库进行了一系列改进。这些更新展示了我们致力于在所有领域开发通用且可扩展 API 的决心，旨在让社区更轻松地基于 PyTorch 构建生态系统项目。

伴随 1.13 版本，我们同时也发布了 PyTorch 各库的更新，详情如下。

TorchAudio

（测试版）混合 Demucs 模型与流水线

混合 Demucs 是一种音乐源分离模型，同时使用了频谱图和时域特征。它在 Sony^® 音乐解混挑战赛中展现了最先进的性能。（引用：https://arxiv.org/abs/2111.03600）

TorchAudio v0.13 版本包含以下功能：

• MUSDB_HQ 数据集，用于混合 Demucs 训练（文档）
• 混合 Demucs 模型架构（文档）
• 三个适用于不同采样率范围的工厂函数
• 预训练流水线（文档）
• 预训练流水线在 MUSDB_HQ 测试集上的 SDR 结果
• 使用预训练流水线进行音乐源分离的教程（文档）

流水线	全部	鼓声	贝斯	其他	人声
HDEMUCS_HIGH_MUSDB*	6.42	7.76	6.51	4.47	6.93
HDEMUCS_HIGH_MUSDB_PLUS**	9.37	11.38	10.53	7.24	8.32

* 在 MUSDB-HQ 数据集的训练数据上进行训练。
** 在 MUSDB-HQ 的训练集和测试集，以及 Meta 内部专门制作的 150 首额外歌曲上进行训练。

from torchaudio.pipelines import HDEMUCS_HIGH_MUSDB_PLUS

bundle = HDEMUCS_HIGH_MUSDB_PLUS
model = bundle.get_model()
sources_list = model.sources

mixture, samplerate = torchaudio.load("song.wav")
sources = model(mixture)
audios = dict(zip(sources_list, sources)

特别感谢 Alexandre Defossez 提供的指导。

（测试版）SUPERB 基准测试的数据集与元数据模式

TorchAudio 增加了对用于自监督学习模型基准测试的各种下游任务音频数据集的支持。随着几个新数据集的加入，现在支持 SUPERB 基准测试 v1 中的下游任务，相关信息可在 s3prl 仓库中找到。

对于这些数据集，我们还通过 get_metadata 函数增加了元数据支持，从而在无需加载波形的情况下实现更快的数据集迭代或预处理。该函数返回与 __getitem__ 相同的特征，但它返回的是相对波形路径，而不是加载后的波形。

具有元数据功能的数据集：

• LIBRISPEECH（文档）
• LibriMix（文档）
• QUESST14（文档）
• SPEECHCOMMANDS（文档）
• （新增）FluentSpeechCommands（文档）
• （新增）Snips（文档）
• （新增）IEMOCAP（文档）
• （新增）VoxCeleb1（识别, 验证）

（测试版）CTC 集束搜索解码中的自定义语言模型支持

TorchAudio 在 0.12 版本中发布了带有 KenLM 语言模型支持的 CTC 集束搜索解码器。在本版本中，我们增加了使用 torchaudio.models.decoder.CTCDecoderLM 包装器创建与该解码器兼容的自定义 Python 语言模型的功能。

有关使用自定义语言模型的更多信息，请参阅文档和教程。

（测试版）StreamWriter

torchaudio.io.StreamWriter 是一个用于编码包括音频和视频在内的媒体的类。它能够处理多种编解码器、分块编码以及 GPU 编码。

writer = StreamWriter("example.mp4")
writer.add_audio_stream(
    sample_rate=16_000,
    num_channels=2,
)
writer.add_video_stream(
    frame_rate=30,
    height=96,
    width=128,
    format="rgb24",
)
with writer.open():
    writer.write_audio_chunk(0, audio)
    writer.write_video_chunk(1, video)

有关更多信息，请参考文档及以下教程：

• StreamWriter 基本用法
• StreamWriter 进阶用法
• 硬件加速视频解码与编码

TorchData

有关变更和新功能的完整列表，请访问我们仓库的 0.5.0 发布说明。

（原型）DataLoader2

DataLoader2 在上一版本中引入，用于执行 DataPipe 图，支持多进程/分布式数据加载的动态分片、多个后端 ReadingServices 以及 DataPipe 图的就地修改（例如 shuffle 控制）。

在本版本中，我们进一步巩固了 DataLoader2 的 API，详细文档请点击此处。我们持续欢迎早期采用者提供反馈和贡献。如果您有兴趣试用，建议安装 TorchData 的 nightly 版本。

（测试版）从云服务提供商处进行数据加载

我们通过 DataPipes 扩展了对从更多云存储提供商处加载数据的支持，现已覆盖 AWS、Google Cloud Storage 和 Azure。同时提供一份教程。我们期待您的反馈和功能需求。

我们还进行了一个简单的基准测试，比较了在 AWS EC2 实例上从 AWS S3 加载数据与从挂载卷加载数据的性能。

torch::deploy (测试版)

torch::deploy 现已进入测试阶段！torch::deploy 是一个适用于 Linux 操作系统的 C++ 库，允许在单个进程中运行多个 Python 解释器。您可以直接运行现有的 eager PyTorch 模型，无需为生产推理用例进行任何更改。亮点包括：

• 现有模型可直接运行，无需修改 Python 代码以支持 tracing。
• 全面支持现有的 Python 环境，包括 C 扩展。
• 在多 GPU 服务环境中进行负载均衡时，无需跨越进程边界。
• 模型权重可以在多个 Python 解释器之间共享。
• 大大改进了安装和设置流程。

torch::deploy::InterpreterManager manager(4);

// access one of the 4 interpreters
auto I = manager.acquireOne();

// run infer from your_model.py
I.global("your_model", "infer")({at::randn({10, 240, 320})});

了解更多信息，请访问此处。

（测试版）CUDA/ROCm/CPU 后端

torch::deploy 现在链接到标准 PyTorch Python 发行版，因此 PyTorch 核心支持的所有加速器（如 CUDA 和 AMD/HIP）均可直接使用。

• 可以通过 pip/conda 正常安装 PyTorch 的任何设备变体。
• https://pytorch.ac.cn/get-started/locally/

（原型）aarch64/arm64 支持

torch::deploy 现在对 aarch64 Linux 系统有了初步支持。

• 我们希望能收集相关反馈，并深入了解 eager PyTorch 模型在 arm 架构上的使用案例。
• 了解更多或分享您的使用案例，请访问 https://github.com/pytorch/multipy/issues/64

TorchEval

（原型）为 PyTorch 引入原生指标支持

TorchEval 是一个为那些希望使用高性能指标来实现机器学习模型评估的用户而构建的库。它还提供了一个易于使用的界面，用于利用相同的工具包构建自定义指标。使用 TorchEval 构建指标，可以轻松运行 torch.distributed 的分布式训练循环。

通过我们的文档了解更多，查看我们的示例，或查看我们的 GitHub 仓库。

TorchMultimodal 发布 (测试版)

请关注 11 月初即将发布的博客，我们将详细介绍 TorchMultimodal。这是一个用于大规模训练最先进（SoTA）多任务多模态模型的 PyTorch 领域库。在此期间，欢迎通过我们的教程体验该库和模型。

TorchRec

（原型）简化的优化器融合 API

我们提供了一个更简化、更直观的 API，通过 apply_optimizer_in_backward 来设置融合优化器。这种新方法允许按每个参数指定优化器设置，分片模块将相应地配置 FBGEMM 的 TableBatchedEmbedding 模块。此外，这让 TorchRec 的规划器能够考量优化器的内存使用情况，从而缓解使用从规划器生成的计划运行 Adam 时出现的内存溢出（OOM）报告。

（原型）简化的分片 API

我们引入了分片（shard）API，现在允许仅对模型中的嵌入模块进行分片，并作为当前主入口点 DistributedModelParallel 的替代方案。这使您能够对模型的其余部分进行更精细的控制，对于自定义并行化逻辑和推理用例（可能不需要在密集层上进行任何并行化）非常有用。我们还引入了 construct_module_sharding_plan，为 TorchRec 分片器提供了更简单的接口。

（测试版）量化通信

在模型并行训练期间，对集体通信调用中的张量应用量化或混合精度，可以极大地提高训练效率，且对模型质量几乎没有影响。TorchRec 现在集成了 FBGEMM GPU 提供的量化通信库，并提供了一个接口来构建围绕分片模块 output_dist 中的 all_to_all 和 reduce_scatter 集体调用的编码器和解码器（编解码器）。我们还允许您构建自己的编解码器以应用于分片模块。FBGEMM 提供的编解码器支持 FP16、BF16、FP8 和 INT8 压缩，并且您可以对前向传播和反向传播使用不同的量化方式。

TorchSnapshot (测试版)

随 PyTorch 1.13 一起，我们发布了 TorchSnapshot 测试版。这是一个为 PyTorch 应用设计的高性能、内存高效的检查点库，专为大型、复杂的分布式工作负载而构建。亮点包括：

• 性能：TorchSnapshot 提供了快速的检查点实现，采用多种优化，包括针对大多数张量类型的零拷贝序列化、重叠的设备到主机拷贝与存储 I/O，以及并行化存储 I/O。
• 内存使用：TorchSnapshot 的内存使用适应主机可用资源，极大地减少了保存和加载检查点时出现内存溢出问题的概率。
• 易用性：简单且在分布式与非分布式工作负载之间保持一致的 API。

通过我们的教程了解更多。

TorchVision

我们很高兴推出 torchvision v0.14 （发布说明）。此版本引入了一个新的模型注册 API，以帮助用户检索和列出模型及权重。它还包括新的图像和视频分类模型，如 MViT、S3D、Swin Transformer V2 和 MaxViT。最后，我们还有新的原语和增强功能，如 PolynomialLR 调度器和 SimpleCopyPaste。

（测试版）模型注册 API

继上一版本发布多权重支持 API 之后，我们添加了一个新的模型注册 API，以帮助用户检索模型和权重。现在在 torchvision.models 模块下有 4 个新方法：get_model、get_model_weights、get_weight 和 list_models。以下是使用它们的示例。

import torchvision
from torchvision.models import get_model, get_model_weights, list_models


max_params = 5000000

tiny_models = []
for model_name in list_models(module=torchvision.models):
    weights_enum = get_model_weights(model_name)
    if len([w for w in weights_enum if w.meta["num_params"] <= max_params]) > 0:
        tiny_models.append(model_name)

print(tiny_models)
# ['mnasnet0_5', 'mnasnet0_75', 'mnasnet1_0', 'mobilenet_v2', ...]

model = get_model(tiny_models[0], weights="DEFAULT")
print(sum(x.numel() for x in model.state_dict().values()))
# 2239188

（测试版）新的视频分类模型

我们添加了两个新的视频分类模型：MViT 和 S3D。MViT 是一种最先进的视频分类 Transformer 模型，在 Kinetics400 数据集上具有 80.757% 的准确率，而 S3D 是一种相对较小但在其规模下具有良好准确率的模型。这些模型的使用方式如下：

import torch
from torchvision.models.video import *

video = torch.rand(3, 32, 800, 600)
model = mvit_v2_s(weights="DEFAULT")
# model = s3d(weights="DEFAULT")
model.eval()
prediction = model(images)

下表显示了在 Kinetics400 数据集上测试的新视频分类模型的准确率：

模型	Acc@1	Acc@5
mvit_v1_b	81.474	95.776
mvit_v2_s	83.196	96.36
s3d	83.582	96.64

我们要感谢 Haoqi Fan、Yanghao Li、Christoph Feichtenhofer 和 Wan-Yen Lo 在 PyTorchVideo 方面的工作，以及他们在 MViT 模型开发期间的支持。我们要感谢 Sophia Zhi 为在 torchvision 中实现 S3D 模型所做的贡献。

（稳定）新架构与模型变体

对于分类模型，我们添加了 Swin Transformer V2 架构及其微型/小型/基础变体的预训练权重。此外，我们还添加了对 MaxViT Transformer 的支持。以下是使用这些模型的示例：

import torch
from torchvision.models import *

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

下表显示了在 ImageNet1K 数据集上测试的模型准确率：

模型	Acc@1	Acc@1 较 V1 的变化	Acc@5	Acc@5 较 V1 的变化
swin_v2_t	82.072	+ 0.598	96.132	+ 0.356
swin_v2_s	83.712	+ 0.516	96.816	+ 0.456
swin_v2_b	84.112	+ 0.530	96.864	+ 0.224
maxvit_t	83.700	–	96.722	–

我们要感谢 Ren Pang 和 Teodor Poncu 为 torchvision 贡献了这两个模型。

（稳定）新原语与增强功能

在此版本中，我们在参考脚本中添加了 SimpleCopyPaste 增强功能，并将 PolynomialLR 调度器上游合并至 PyTorch 核心。我们要感谢 Lezwon Castelino 和 Federico Pozzi 的贡献。我们正继续努力通过社区的帮助，添加更多最先进的原语、增强功能和架构来现代化 TorchVision。如果您有兴趣做出贡献，请查看以下议题。

Torch-TensorRT

（原型）带有 FX2TRT 前端的 TensorRT

Torch-TensorRT 是 PyTorch 与 TensorRT 的集成，在 NVIDIA GPU 上提供高性能推理。Torch-TRT 允许直接在 PyTorch 中优化模型以进行部署，提供高达 6 倍的性能提升。

Torch-TRT 是一个提前（AoT）编译器，它接收 nn.Module 或 TorchScript 模块，在 TensorRT 中优化兼容的子图，并将其余部分留给 PyTorch 运行。这为用户提供了 TensorRT 的高性能，以及 Torch 的可用性和熟悉度。

Torch-TensorRT 是 PyTorch 生态系统的一部分，于 2021 年 11 月发布 v1.0 版本。目前有两个不同的前端：Torchscript 和 FX。每个前端都提供相同的价值主张和底层操作，主要区别在于输入和输出格式（TS 与 FX / Python）。

Torchscript 前端已包含在 v1.0 中，被认为是稳定的。FX 前端首次发布于 v1.2，被视为测试版。

相关链接：

• GitHub
• 文档
• 通用 (TS) 入门指南
• FX 入门指南

（稳定）引入 Torch-TensorRT

Torch-TensorRT 是一个 PyTorch 集成，它利用 TensorRT 在 NVIDIA GPU 上的推理优化。它利用 TensorRT 的优化（如 FP16 和 INT8 降低精度、图优化、算子融合等），同时在 TensorRT 不支持模型子图时提供回退到原生 PyTorch 的功能。目前库中存在两个前端路径，有助于将 PyTorch 模型转换为 TensorRT 引擎。一条路径是通过 Torch Script (TS)，另一条是通过 FX 前端。即模型被 TS 或 FX 跟踪到其 IR 图中，然后从中转换为 TensorRT。

通过我们的教程了解更多。

TorchX

TorchX 0.3 更新包括一个新的 list API、实验追踪、弹性训练和改进的调度器支持。还有一个使用 TorchX + Ax 的新多目标 NAS 教程。

（原型）List

新添加的 list 命令和 API 允许您直接在 TorchX 中为给定的调度器列出最近启动的作业及其状态。

• 这消除了使用辅助工具来列出作业的需求。
• 完全程序化地访问最近的作业，以便与自定义工具集成。

$ torchx list -s kubernetes
APP HANDLE                                                       APP STATUS
-----------------------------------------------            -----------------
kubernetes://torchx/default:train-f2nx4459p5crr   SUCCEEDED

通过我们的文档了解更多。

（原型）Tracker

TorchX Tracker 是一个新的原型库，它提供了一个灵活且可定制的实验和工件追踪界面。这允许您跨多个步骤追踪作业的输入和输出，从而更轻松地将 TorchX 与流水线及其他外部系统一起使用。

from torchx import tracker

app_run = tracker.app_run_from_env()
app_run.add_metadata(lr=lr, gamma=gamma) # hyper parameters
app_run.add_artifact("model", "storage://path/mnist_cnn.pt") # logs / checkpoints
app_run.add_source(parent_run_id, "model") # lineage

例子

• https://github.com/pytorch/torchx/tree/main/torchx/examples/apps/tracker
• https://pytorch.ac.cn/torchx/main/tracker.html

（原型）弹性训练与自动缩放

Ray 和 Kubernetes 上的弹性——在使用受支持的调度器时，自动扩展分布式训练作业。通过我们的文档了解更多。

（原型）调度器改进：IBM® Spectrum LSF

增加了对 IBM Spectrum LSF 调度器的原型支持。

（测试版）AWS Batch 调度器

AWS Batch 调度器集成现已进入测试阶段。

• 现在支持日志获取和作业列表显示。
• 增加了作业优先级和队列策略的配置。
• 通过 ui_url 轻松访问作业 UI：https://pytorch.ac.cn/torchx/main/schedulers/aws_batch.html

（原型）AnyPrecision 优化器

AdamW 优化器的直接替代品，可减少 GPU 内存，实现两个主要功能：

• 能够以完全 BFloat16 成功训练整个模型流水线。Kahan 求和确保了精度。这可以提高训练吞吐量，特别是在超大模型上，通过减少内存和提高计算速度实现。
• 能够将方差状态更改为 BFloat16。这可以减少模型训练所需的总内存，并带来额外的速度提升。

在此处查找更多信息。