PyTorch 1.13 中的新库更新

摘要

除了 PyTorch 1.13 版本之外，我们还对当前的 PyTorch 库进行了一些改进。这些更新表明我们专注于在所有领域开发通用且可扩展的 API，以便我们的社区更容易在 PyTorch 上构建生态系统项目。

除了 **1.13**，我们还将发布 PyTorch 库的更新，请在下面找到它们。

TorchAudio

（Beta）Hybrid Demucs 模型和流水线

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

TorchAudio v0.13 版本包括以下功能

• MUSDB_HQ 数据集，用于 Hybrid Demucs 训练（文档）
• Hybrid 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 的指导。

（Beta）SUPERB 基准测试的数据集和元数据模式

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

对于这些数据集，我们还通过 `get_metadata` 函数添加了元数据支持，从而无需加载波形即可实现更快的数据集迭代或预处理。该函数返回与 `__getitem__` 相同的功能，只是它返回相对波形路径而不是加载的波形。

具有元数据功能的数据集

• LIBRISPEECH (文档)
• LibriMix (文档)
• QUESST14 (文档)
• SPEECHCOMMANDS (文档)
• （新）FluentSpeechCommands (文档)
• （新）Snips (文档)
• （新）IEMOCAP (文档)
• （新）VoxCeleb1 (识别、验证)

（Beta）CTC 波束搜索解码中的自定义语言模型支持

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

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

（Beta）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 的每夜版本。

（Beta）从云服务提供商加载数据

我们扩展了通过 DataPipes 从其他云存储提供商加载数据的支持，现在涵盖 AWS、Google Cloud Storage 和 Azure。还提供了教程。我们欢迎反馈和功能请求。

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

torch::deploy (Beta)

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

• 现有模型开箱即用——无需修改您的 python 代码即可支持跟踪。
• 完全支持您现有的 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})});

在此处了解更多信息：https://github.com/pytorch/multipy。

（Beta）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 发布（Beta）

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

TorchRec

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

我们提供了一个简化且更直观的 API，用于通过 apply_optimizer_in_backward 设置融合优化器设置。这种新方法能够以每参数为基础指定优化器设置，并且分片模块将相应地配置 FBGEMM 的 TableBatchedEmbedding 模块。此外，这现在允许 TorchRec 的规划器考虑优化器内存使用。这应该可以缓解使用规划器生成的计划使用 Adam 后分片作业 OOM 的报告。

（原型）简化的分片 API

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

（Beta）量化通信

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

TorchSnapshot (Beta)

与 PyTorch 1.13 一起，我们发布了 TorchSnapshot 的 Beta 版本，它是一个高性能、内存高效的 PyTorch 应用程序检查点库，旨在考虑大型、复杂的分布式工作负载。亮点包括

• 性能：TorchSnapshot 提供了一种快速的检查点实现，采用了各种优化，包括大多数张量类型的零拷贝序列化、重叠的设备到主机复制和存储 I/O、并行存储 I/O
• 内存使用：TorchSnapshot 的内存使用量适应主机的可用资源，大大降低了保存和加载检查点时出现内存不足问题的可能性
• 可用性：分布式和非分布式工作负载之间一致的简单 API

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

TorchVision

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

（Beta）模型注册 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

（Beta）新视频分类模型

我们添加了两个新的视频分类模型，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 架构以及其 tiny/small/base 变体的预训练权重。此外，我们还添加了对 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 Core。我们要感谢 Lezwon Castelino 和 Federico Pozzi 的贡献。我们正在继续努力，在社区的帮助下通过添加更多 SoTA 基元、增强功能和架构来现代化 TorchVision。如果您有兴趣做出贡献，请查看以下问题。

Torch-TensorRT

（原型）带有 FX2TRT 前端的 TensorRT

Torch-TensorRT 是 TensorRT 的 PyTorch 集成，可在 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 中发布，应视为 Beta 版。

相关链接

• 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 图中，然后从 IR 图转换为 TensorRT。

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

TorchX

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

（原型）列表

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

• 这消除了使用辅助工具来列出作业的需要。
• 对最近作业的完全编程访问，用于与自定义工具集成。

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

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

（原型）跟踪器

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 调度器的原型支持。

（Beta）AWS Batch 调度器

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

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

（原型）AnyPrecision 优化器

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

• 能够在完整的 BFloat16 中成功训练整个模型管道。Kahan 求和确保精度。这可以通过减少内存和提高计算速度来提高训练吞吐量，尤其是在大型模型上。
• 能够将方差状态更改为 BFloat16。这可以减少模型训练所需的总内存，并提高速度。

在此处查找更多信息：https://github.com/pytorch/torchdistx/pull/52。