PyTorch 2.2 新版图书馆更新

由 PyTorch 团队

摘要

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

*查看先前版本或（不稳定）夜间版本，请点击顶部左侧菜单中的“搜索文档”旁边的版本。

TorchRL

特性：TorchRL 的离线 RL 数据中心

TorchRL 现在提供最大的离线 RL 和模仿学习数据集之一，所有数据都采用单一数据格式（TED，即 TorchRL Episode Data 格式）。这使得在单个训练循环中轻松交换不同来源的数据成为可能。现在还可以通过 ReplayBufferEnsemble 类轻松组合不同来源的数据集。数据处理完全可定制。数据来源包括模拟任务（Minari、D4RL、VD4RL）、机器人数据集（Roboset、OpenX Embodied 数据集）和游戏（GenDGRL/ProcGen、Atari/DQN）。请查看文档中的相关信息。

除了这些变化之外，我们的重放缓冲区现在可以使用 .dumps() 方法在磁盘上进行转储，该方法将使用 TensorDict API 在磁盘上序列化缓冲区，比使用 torch.save 更快、更安全、更高效。

最后，重放缓冲区现在可以在同一台机器上的不同进程中读取和写入，无需用户编写任何额外代码！

TorchRL2Gym 环境 API

为了便于 TorchRL 与现有代码库的集成，并享受 TorchRL 环境 API 的所有功能（设备上的执行、批量操作、转换等），我们提供了一个 TorchRL-to-gym API，允许用户将任何他们想要的 gym 或 gymnasium 环境注册到其中。这反过来可以将 TorchRL 变成一个通用的 lib-to-gym 转换器，适用于无状态（例如，dm_control）和有状态（Brax、Jumanji）环境。该功能在文档中有详细说明。info_dict 读取 API 也得到了改进。

环境加速

我们添加了在 ParallelEnv 中在不同于数据交付环境的其他环境中执行环境的选项。我们还加速了 GymLikeEnv 类，使其现在与 gym 本身具有竞争力。

规模化目标

最受欢迎的 RLHF 和大规模训练（PPO 和 A2C）的目标现在与 FSDP 和 DDP 模型兼容！

TensorDict

功能：MemoryMappedTensor 替代 MemmapTensor

我们提供了一个更高效的 mmap 后端，用于 TensorDict；MemoryMappedTensor，它直接继承自 torch.Tensor。它附带了一系列构建工具，如 from_tensor ， empty 等等。MemoryMappedTensor 现在比其对应版本更加安全且速度更快。该库与之前的类保持完全兼容，以方便过渡。

我们还引入了一套新的多线程序列化方法，使得 tensordict 的序列化在速度上与 torch.save 具有高度竞争力，序列化和反序列化的速度比 torch.save 快 3 倍以上。

功能：TensorDict 中的非张量数据

通过 NonTensorData tensorclass 无法传输非张量数据。这使得构建带有元数据的 tensordicts 成为可能。 memmap -API 与这些值完全兼容，允许用户无缝序列化和反序列化此类对象。要在 tensordict 中存储非张量数据，只需使用 __setitem__ 方法进行赋值。

效率提升

改进了多个运行时方法，例如 unbind、split、map 甚至 TensorDict 的实例化。查看我们的基准测试！

TorchRec/fbgemm_gpu

VBE

TorchRec 现已原生支持 EmbeddingBagCollection 模块中的 VBE（可变批量嵌入）。这允许每个特征具有可变的批量大小，从而解锁稀疏输入数据的去重，可以大大加快嵌入查找和全对全的时间。要启用，只需使用 stride_per_key_per_rank 和 inverse_indices 字段初始化 KeyedJaggedTensor ，分别指定每个特征的批量大小和重新索引嵌入输出的逆索引即可。

除了 TorchRec 库的更改外，fbgemm_gpu 还增加了对 TBE 中每个特征可变批量大小的支持。在分割 TBE 训练中，对于加权和未加权情况均启用了 VBE。要使用 VBE，请确保使用最新的 fbgemm_gpu 版本。

嵌入卸载

该技术指的是使用 CUDA UVM 来缓存“热点”嵌入（即在主机内存上存储嵌入表，在 HBM 内存上缓存），并预取缓存。嵌入卸载允许使用更少的 GPU 运行更大的模型，同时保持有竞争力的性能。请使用预取流水线（PrefetchTrainPipelineSparseDist）并通过规划器中的约束传递每个表的缓存加载因子和 prefetch_pipeline 标志来使用此功能。

Fbgemm_gpu 在 v0.5.0 版本中为 TBE 性能加速引入了 UVM 缓存流水线预取。这允许缓存插入与 TBE 正向/反向操作并行执行。要启用此功能，请确保使用最新的 fbgemm_gpu 版本。

Trec.shard/shard_modules

这些 API 用其分片变体替换嵌入子模块。shard API 适用于单个嵌入模块，而 shard_modules API 替换所有嵌入模块，不会触及其他非嵌入子模块。

嵌入式分片的行为与之前的 TorchRec DistributedModuleParallel 行为相似，但 ShardedModules 已被设计为可组合的，这意味着模块由 TableBatchedEmbeddingSlices 支持，这些是底层 TBE（包括.gradle）的视图。这意味着现在融合的参数可以通过 named_parameters()返回，包括在 DistributedModuleParallel 中。

TorchVision

V2 转换现在已稳定！

直到如今， torchvision.transforms.v2 命名空间仍处于 BETA 阶段。现在它已经稳定了！无论你是 Torchvision 转换的新手，还是已经熟悉它们，我们都鼓励你从“使用 v2 转换入门”开始，以了解更多关于新 v2 转换可以做什么。

浏览我们的主要文档以获取一般信息和性能技巧。可用的转换和功能在 API 参考中列出。更多信息和学习教程也可以在我们的示例画廊中找到，例如“Transforms v2：端到端目标检测/分割示例”或“如何编写自己的 v2 转换”。

向 torch.compile() 支持迈进

我们正在逐步将 torch.compile() 支持添加到 torchvision 接口，减少图断裂并允许动态形状。

torchvision 操作符（ nms 、 [ps_]roi_align 、 [ps_]roi_pool 和 deform_conv_2d ）现在与 torch.compile 和动态形状兼容。

在转换方面，大多数底层内核（如 resize_image() 或 crop_image() ）应能正确编译，没有图断裂并支持动态形状。我们仍在解决剩余的边缘情况，逐步向全面功能支持和类迈进，您应该期待在下一个版本中看到这方面的更多进展。