PyTorch 1.13 中的新库更新

摘要

我们正在为当前的 PyTorch 库带来一系列改进，同时发布 PyTorch 1.13 版本。这些更新表明我们专注于开发跨所有领域的通用和可扩展 API，以便我们的社区更容易在 PyTorch 上构建生态系统项目。

随着 1.13 的发布，我们正在发布 PyTorch 库的更新，请在下面找到它们。

TorchAudio

（测试版）Hybrid Demucs 模型和管道

Hybrid Demucs 是一种音乐源分离模型，它使用频谱图和时域特征。它在 Sony^® Music DeMixing Challenge 中展示了最先进的性能。（引用： 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 的指导。

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

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

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

具有元数据功能的数据集

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

（测试版）CTC Beam Search 解码中对自定义语言模型的支持

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

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

（测试版）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 S3 和 AWS EC2 实例上的附加卷加载数据的性能。

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})});

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

（测试版）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 架构以及其 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	相对于 V1 的 Acc@1 变化	Acc@5	相对于 V1 的 Acc@5 变化
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 生态系统的一部分，并于 21 年 11 月发布了 v1.0。目前有两个不同的前端：Torchscript 和 FX。每个都提供相同的价值主张和底层操作，主要区别在于输入和输出格式（TS vs 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 图中，然后从 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

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

（原型）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。这可以减少模型训练所需的总内存，并提高速度。

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