跳转到主要内容
博客

PyTorch 2.0 中的新库更新

作者： 2023年3月15日2024年11月14日暂无评论

摘要

随着 PyTorch 2.0 发布，我们正在对当前的 PyTorch 库进行一系列改进。这些更新表明我们专注于在所有领域开发通用且可扩展的 API，以便社区更轻松地在 PyTorch 上构建生态系统项目。

与 2.0 一起，我们还将发布 PyTorch 领域库的一系列 Beta 更新，包括在树内的库和单独的库，如 TorchAudio、TorchVision 和 TorchText。TorchX 的更新也正在发布，因为它将转为社区支持模式。请在下方找到最新稳定版本和更新列表。

最新稳定库版本 (完整列表)

TorchArrow 0.1.0TorchRec 0.4.0TorchVision 0.15
TorchAudio 2.0TorchServe 0.7.1TorchX 0.4.0
TorchData 0.6.0TorchText 0.15.0PyTorch on XLA Devices 1.14

*要查看 以前的版本 或（不稳定）的夜间版本，请点击“搜索文档”上方左上角菜单中的版本。

TorchAudio

[Beta] 数据增强操作

此版本在 torchaudio.functional 和 torchaudio.transforms 下添加了多个数据增强操作

  • torchaudio.functional.add_noise
  • torchaudio.functional.convolve
  • torchaudio.functional.deemphasis
  • torchaudio.functional.fftconvolve
  • torchaudio.functional.preemphasis
  • torchaudio.functional.speed
  • torchaudio.transforms.AddNoise
  • torchaudio.transforms.Convolve
  • torchaudio.transforms.Deemphasis
  • torchaudio.transforms.FFTConvolve
  • torchaudio.transforms.Preemphasis
  • torchaudio.transforms.Speed
  • torchaudio.transforms.SpeedPerturbation

这些操作可用于合成多样化训练数据，以提高下游模型的泛化能力。

有关使用详情，请参阅 functional 和 transform 文档以及 音频数据增强 教程。

[Beta] WavLM 和 XLS-R 模型

此版本添加了两个用于语音和音频的自监督学习模型。

  • WavLM，对噪声和混响具有鲁棒性。
  • XLS-R，在跨语言数据集上进行训练。

除了模型架构，torchaudio 还支持相应的预训练管道

  • torchaudio.pipelines.WAVLM_BASE
  • torchaudio.pipelines.WAVLM_BASE_PLUS
  • torchaudio.pipelines.WAVLM_LARGE
  • torchaudio.pipelines.WAV2VEC_XLSR_300M
  • torchaudio.pipelines.WAV2VEC_XLSR_1B
  • torchaudio.pipelines.WAV2VEC_XLSR_2B

有关使用详情，请参阅 工厂函数 和 预训练管道 文档。

TorchRL

torchrl 的初始版本包含涵盖整个强化学习领域的多个功能。TorchRL 已经可以在在线、离线、多智能体、多任务和分布式强化学习设置等中使用。详见下方

[Beta] 环境封装器和变换

torchrl.envs 包含围绕常见环境库的多个封装器。这允许用户轻松地将一个库替换为另一个库。这些封装器在这些模拟器和 torchrl 之间构建了一个接口

  • dm_control
  • Gym
  • Brax
  • EnvPool
  • Jumanji
  • Habitat

它还附带了许多常用的变换和矢量化环境实用程序，允许在模拟库中快速执行。请参阅 文档 以获取更多详细信息。

[Beta] 数据收集器

通过使用单进程或多进程/分布式数据收集器，可以在强化学习中轻松进行数据收集，这些收集器在环境中按所需持续时间执行策略，并根据用户需求提供样本。这些可以在 torchrl.collectors 中找到，并在此处 记录。

[Beta] 目标模块

torchrl.objectives 中包含多个目标函数，其中包括

  • 通用的 PPOLoss 类和派生的 ClipPPOLoss 和 KLPPOLoss
  • SACLoss 和 DiscreteSACLoss
  • DDPGLoss
  • DQNLoss
  • REDQLoss
  • A2CLoss
  • TD3Loss
  • ReinforceLoss
  • Dreamer

矢量化值函数操作符也出现在库中。在此处 查看文档。

[Beta] 模型和探索策略

我们提供多种模型、模块和探索策略。在 文档 中获取详细描述。

[Beta] 可组合回放缓冲区

提供了一个可组合的回放缓冲区类，可用于在多个上下文（包括单智能体和多智能体、在线和离线策略等）中存储数据。组件包括

  • 存储 (列表、物理或基于内存的连续存储)
  • 采样器 (优先、无重复采样器)
  • 写入器
  • 添加变换的可能性

回放缓冲区和其他数据实用程序在此处 记录。

[Beta] 日志工具和训练器

我们支持多种日志工具，包括 tensorboard、wandb 和 mlflow。

我们提供了一个通用的 Trainer 类，可以轻松进行代码重用和检查点。

这些功能在此处 记录。

TensorDict

TensorDict 是 PyTorch 的新数据载体。

[Beta] TensorDict：PyTorch 专用字典

TensorDict 允许您对单个容器承载的张量批次执行许多常见操作。TensorDict 支持多种形状和设备或存储操作，并且可以轻松地用于分布式设置。查阅 文档 以了解更多信息。

[Beta] @tensorclass：PyTorch 的数据类

与 TensorDict 类似，tensorclass 提供了编写具有内置 torch 功能（如形状或设备操作）的数据类的机会。

[Beta] tensordict.nn：TensorDict 专用模块

tensordict.nn 模块 提供了专门的 nn.Module 子类，使构建可使用 TensorDict 输入执行的任意复杂图变得容易。它与最新的 PyTorch 功能兼容，如 functorch、torch.fx 和 torch.compile。

TorchRec

[Beta] KeyedJaggedTensor All-to-All 重新设计和输入分布式融合

我们观察到由于具有多个大型 KJT 需要重新分配的模型中稀疏数据分布瓶颈导致的性能下降。

为了解决这个问题，我们改变了通信模式，以在初始集合中传输支持实际 KJT 张量数据的集合调用所需的最少数据。在初始集合中发送的这些“拆分”数据意味着在通信流中传输更多数据，但 CPU 阻塞的时间显著缩短，从而带来更好的整体 QPS。

此外，我们修改了 TorchRec 训练管道，将拆分的初始集合调用分组在一起，然后再启动更昂贵的 KJT 张量集合调用。这种融合最大限度地减少了 CPU 阻塞时间，因为启动每个后续输入分布不再依赖于先前的输入分布。

通过此功能，现在可以在不同等级上原生支持可变批次大小。这些功能在此处 记录。

TorchVision

[Beta] 将 TorchVision 的变换扩展到对象检测、分割和视频任务

TorchVision 正在扩展其 Transforms API！新功能如下：

  • 您可以将它们不仅用于图像分类，还用于目标检测、实例和语义分割以及视频分类。
  • 您可以使用新的函数式转换来转换视频、边界框和分割掩码。

从 我们的文档 中了解有关这些新变换的更多信息，并在 我们的专用问题 中提交任何反馈。

TorchText

[Beta] 为 TorchText 库添加可脚本化的 T5 和 Flan-T5，并支持增量解码！

TorchText 增加了 T5 模型架构，并提供了 原始 T5 论文 和 Flan-T5 的预训练权重。该模型完全可脚本化，并具有优化的 多头注意力实现。我们提供了几个如何利用该模型的示例，包括摘要、分类和翻译。

有关更多详细信息，请参阅 我们的文档。

TorchX

TorchX 正在转为社区支持模式。更多详细信息将在稍后公布。