总结
在 PyTorch 2.2 发布的同时，我们为当前的 PyTorch 库带来了一系列改进。这些更新表明我们致力于在所有领域开发通用且可扩展的 API，以便我们的社区更容易在 PyTorch 上构建生态系统项目。
| 最新的稳定库版本 (完整列表)* | ||
| TorchArrow 0.1.0 | TorchRec 0.6.0 | TorchVision 0.17 |
| TorchAudio 2.2.0 | TorchServe 0.9.0 | TorchX 0.7.0 |
| TorchData 0.7.1 | TorchText 0.17.0 | PyTorch on XLA Devices 2.1 |
*要查看 以前的版本 或（不稳定）的夜间版本，请点击“搜索文档”上方左上角菜单中的版本。
TorchRL
功能：TorchRL 的离线强化学习数据中心
TorchRL 现在提供了一个最大的离线强化学习和模仿学习数据集中心，所有数据都采用单一数据格式（TED，即 TorchRL Episode Data 格式）。这使得在单个训练循环中轻松切换不同来源成为可能。现在还可以通过 ReplayBufferEnsemble 类轻松组合来自不同来源的数据集。数据处理完全可定制。来源包括模拟任务（Minari、D4RL、VD4RL）、机器人数据集（Roboset、OpenX Embodied 数据集）和游戏（GenDGRL/ProcGen、Atari/DQN）。请在文档中查看。
除了这些更改之外，我们的重放缓冲区现在可以使用 .dumps() 方法转储到磁盘上，该方法将使用 TensorDict API 将缓冲区序列化到磁盘上，该 API 比使用 torch.save 更快、更安全、更高效。
最后，重放缓冲区现在可以在同一台机器上的不同进程中读取和写入，而无需用户编写任何额外的代码！
TorchRL2Gym 环境 API
为了促进 TorchRL 在现有代码库中的集成并享受 TorchRL 环境 API 的所有功能（设备上的执行、批处理操作、变换等），我们提供了一个 TorchRL 到 gym 的 API，允许用户在 gym 或 gymnasium 中注册他们想要的任何环境。这反过来可以使 TorchRL 成为一个通用的库到 gym 转换器，适用于无状态（例如，dm_control）和有状态（Brax、Jumanji）环境。该功能在文档中详细说明。info_dict 读取 API 也得到了改进。
环境加速
我们增加了在 ParallelEnv 中在与用于提供数据的环境不同的环境上执行环境的选项。我们还加快了 GymLikeEnv 类的速度，使其现在可以与 gym 本身竞争。
缩放目标
最流行的 RLHF 和大规模训练目标（PPO 和 A2C）现在与 FSDP 和 DDP 模型兼容！
TensorDict
功能：MemoryMappedTensor 取代 MemmapTensor
我们为 TensorDict 提供了一个更高效的 mmap 后端；MemoryMappedTensor，它直接继承自 torch.Tensor。它附带了一系列用于构建的实用程序，例如 from_tensor、empty 等。MemoryMappedTensor 现在比其对应物更安全、更快。该库与以前的类完全兼容，以方便过渡。
我们还引入了一组新的多线程序列化方法，使 tensordict 序列化与 torch.save 极具竞争力，LLM 的序列化和反序列化速度比使用 torch.save 快 3 倍以上。
功能：TensorDict 中的非张量数据
现在可以通过 NonTensorData 张量类携带非张量数据。这使得可以构建带有元数据的 tensordict。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 中添加了对每个特征可变批次大小的支持。VBE 在加权和未加权情况下都启用了拆分 TBE 训练。要使用 VBE，请务必使用最新的 fbgemm_gpu 版本。
嵌入卸载
此技术是指使用 CUDA UVM 缓存“热”嵌入（即在主机内存中存储嵌入表并在 HBM 内存中进行缓存），并预取缓存。嵌入卸载允许使用更少的 GPU 运行更大的模型，同时保持具有竞争力的性能。使用预取管道（PrefetchTrainPipelineSparseDist）并在规划器中通过约束条件传入每个表的缓存加载因子和预取管道标志以使用此功能。
Fbgemm_gpu 在 v0.5.0 中引入了 UVM 缓存管道预取，以提高 TBE 性能。这允许缓存插入与 TBE 正向/反向并行执行。要启用此功能，请务必使用最新的 fbgemm_gpu 版本。
Trec.shard/shard_modules
这些 API 将嵌入子模块替换为其分片变体。shard API 适用于单个嵌入模块，而 shard_modules API 替换所有嵌入模块，并且不会触及其他非嵌入子模块。
嵌入分片遵循与之前的 TorchRec DistributedModuleParallel 行为相似的行为，不同之处在于 ShardedModules 已变为可组合的，这意味着这些模块由 TableBatchedEmbeddingSlices 支持，这些是底层 TBE（包括 .grad）的视图。这意味着现在使用 named_parameters() 返回融合参数，包括在 DistributedModuleParallel 中。
TorchVision
V2 转换现在稳定了！
torchvision.transforms.v2 命名空间直到现在仍处于 BETA 阶段。它现在稳定了！无论您是 Torchvision 转换的新手，还是已经有经验，我们都鼓励您从V2 转换入门开始，以了解新 V2 转换可以做什么。
浏览我们的主文档以获取一般信息和性能提示。可用的转换和函数在API 参考中列出。您还可以在我们的示例库中找到更多信息和教程，例如V2 转换：端到端目标检测/分割示例或如何编写您自己的 V2 转换。
迈向 torch.compile() 支持
我们正在逐步将 torch.compile() 支持添加到 torchvision 接口，减少图中断并允许动态形状。
torchvision ops（nms、[ps_]roi_align、[ps_]roi_pool 和 deform_conv_2d）现在与 torch.compile 和动态形状兼容。
在转换方面，大部分低级内核（如 resize_image() 或 crop_image()）应该能够正确编译，没有图中断和动态形状。我们仍在解决其余的边缘情况，朝着完整的函数支持和类迈进，您应该会在下一个版本中看到这方面的更多进展。