跳转到主要内容
博客

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 文档以及 Audio Data Augmentation 教程。

[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 的初始版本包括涵盖整个 RL 领域的多种功能。TorchRL 已经可以用于在线、离线、多智能体、多任务和分布式 RL 设置等。请参见下方

[Beta] 环境包装器和转换

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

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

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

[Beta] 数据收集器

通过使用单进程或多进程/分布式数据收集器，可以在所需的持续时间内在环境中执行策略并根据用户需求提供样本，从而简化 RL 中的数据收集。这些可以在 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 提供了使用内置 PyTorch 功能（如形状或设备操作）编写数据类的机会。

[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 正在进入社区支持模式。更多细节将在稍后公布。