由 PyTorch 团队

摘要

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

除了 1.13 版本外,我们还发布了 PyTorch 库的更新,请见下文。

TorchAudio

(Beta) 混合 Demucs 模型和管道

混合 Demucs 是一种同时使用频谱图和时间域特征的音频分离模型。它在索尼 ® 音乐去混响挑战赛中展现了最先进的性能。(参考文献: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 的指导。

(Beta) 为 SUPERB 基准测试添加数据集和元数据模式。

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

对于这些数据集,我们还通过 get_metadata 函数添加了元数据支持,使得在无需加载波形的情况下,可以更快地进行数据集迭代或预处理。该函数返回与 __getitem__ 相同的特征,但返回的是相对波形路径,而不是加载的波形。

具有元数据功能的数据集

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

(Beta) CTC Beam Search 解码中的自定义语言模型支持

TorchAudio 在 0.12 版本中发布了一个 CTC Beam Search 解码器,支持 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)

如需更多信息,请参阅文档和以下教程。

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

在这里了解更多。

(Beta) CUDA/ROCm/CPU 后端

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

(原型)支持 aarch64/arm64 架构。

torch::deploy 现在对 aarch64 Linux 系统有基本支持。

  • 我们正在寻求对此收集反馈,并了解对急切 PyTorch 模型的 arm 使用案例。
  • 在 https://github.com/pytorch/multipy/issues/64 了解更多/分享您的用例

TorchEval

(原型) 为 PyTorch 引入原生指标支持

TorchEval 是一个为希望获得常用指标的高性能实现并评估机器学习模型的用户构建的库。它还提供了一个易于使用的接口,使用相同的工具包构建自定义指标。使用 TorchEval 构建您的指标,使运行带有 torch.distributed 的分布式训练循环变得轻而易举。

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

TorchMultimodal 发布(Beta)

请关注 11 月初即将推出的博客,将详细介绍 TorchMultimodal,这是一个 PyTorch 领域库,用于大规模训练 SoTA 多任务多模态模型;在此期间,请通过我们的教程尝试使用库和模型。

TorchRec

(原型)简化优化器融合 API

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

(原型)简化分片 API

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

(测试版)量化通信

在模型并行训练中对张量进行量化或混合精度调用,可以大大提高训练效率,对模型质量的影响微乎其微。TorchRec 现已集成 FBGEMM GPU 提供的量化通信库,并提供了一个接口来构建围绕 sharded 模块的 output_dist 中的 all_to_all 和 reduce_scatter 集体调用的编码器和解码器(codec)。我们还允许您构建自己的 codec 以应用于您的 sharded 模块。FBGEMM 提供的 codec 支持 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。最后但同样重要的是,我们还引入了新的原语和增强,如多项式 LR 调度器和简单复制粘贴。

(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 数据集上测试的新视频分类模型准确率的表格。

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

我们要感谢范浩奇、李阳豪、克里斯托夫·费希滕霍费尔和罗万恩对 PyTorchVideo 的工作以及他们在 MViT 模型开发过程中的支持。我们还要感谢智晓对在 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 数据集上测试的模型准确率的表格。

Model 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 贡献了 2 个模型。

(稳定) 新的原始函数和增强

在本次版本中,我们在参考脚本中添加了 SimpleCopyPaste 增强,并将 PolynomialLR 调度器升级到 PyTorch Core。我们感谢 Lezwon Castelino 和 Federico Pozzi 的贡献。我们正在继续努力通过社区的帮助,使 TorchVision 现代化,添加更多 SOTA 原始函数、增强和架构。如果您有兴趣贡献,请查看以下问题。

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 中发布,应被视为 Beta 版本。

相关链接:

(稳定版) 推出 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 更新包括新的列表 API、实验跟踪、弹性训练和改进的调度器支持。还提供了一个使用 TorchX + Ax 的新多目标 NAS 教程。

(原型)列表

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

  • 这样就无需使用其他工具来列出作业。
  • 完全程序化访问最近的工作,以便与自定义工具集成。
$ torchx list -s kubernetes
APP HANDLE                                                       APP STATUS
-----------------------------------------------            -----------------
kubernetes://torchx/default:train-f2nx4459p5crr   SUCCEEDED

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

(原型)跟踪器

TorchX 跟踪器是一个新的原型库,它提供了一个灵活且可定制的实验和工件跟踪界面。这允许您跟踪多个步骤中作业的输入和输出,以便更容易地将 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

例如:

(原型)弹性训练和自动扩展

在使用受支持的调度器时,Ray 和 Kubernetes 上的弹性 - 当分布式训练作业自动扩展。了解更多信息,请参阅我们的文档。

(原型)调度器改进:IBM® Spectrum LSF

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

(Beta) AWS Batch 调度器

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

  • 现已支持日志获取和作业列表。
  • 添加了作业优先级和队列策略配置。
  • 通过 ui_url https://maskerprc.github.io/torchx/main/schedulers/aws_batch.html 轻松访问工作 UI

(原型) AnyPrecision 优化器

替换 AdamW 优化器的即插即用方案,可减少 GPU 内存,实现两大主要功能:

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

更多信息请在此处查找。