总结
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 dataset）和游戏（GenDGRL/ProcGen、Atari/DQN）。请在 文档中查看这些内容。
除了这些更改之外，我们的回放缓冲区现在可以使用 .dumps() 方法转储到磁盘上，该方法将使用 TensorDict API 将缓冲区序列化到磁盘上，这比使用 torch.save 更快、更安全、更高效。
最后，现在无需用户额外代码，即可在同一机器上的不同进程中读取和写入回放缓冲区！
TorchRL2Gym 环境 API
为了方便 TorchRL 在现有代码库中的集成并享受 TorchRL 环境 API 的所有功能（设备执行、批处理操作、转换……），我们提供了一个 TorchRL-to-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 tensorclass 携带非张量数据。这使得构建带有元数据的 tensordict 成为可能。memmap-API 与这些值完全兼容，允许用户无缝序列化和反序列化此类对象。要将非张量数据存储在 tensordict 中，只需使用 __setitem__ 方法进行赋值即可。
效率提升
许多方法的运行时间都得到了改进，例如 unbind、split、map 甚至 TensorDict 实例化。查看我们的 基准测试！
TorchRec/fbgemm_gpu
VBE
TorchRec 现在在 EmbeddingBagCollection 模块中原生支持 VBE（可变批处理嵌入）。这允许每个特征的可变批大小，从而实现稀疏输入数据去重，这可以大大加快嵌入查找和 all-to-all 时间。要启用，只需使用 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 转换的新手，还是已经有经验，我们都鼓励您从 使用 transforms v2 入门开始，以了解新 v2 转换的功能。
浏览我们的主要文档以获取一般信息和性能提示。可用的转换和函数在API 参考中列出。您还可以在我们的示例库中找到其他信息和教程，例如Transforms v2：端到端对象检测/分割示例或如何编写自己的 v2 转换。
朝向 torch.compile() 支持
我们正在逐步将 torch.compile() 支持添加到 torchvision 接口中，减少图中断并允许动态形状。
torchvision ops（nms、[ps_]roi_align、[ps_]roi_pool 和 deform_conv_2d）现在与 torch.compile 和动态形状兼容。
在 transforms 方面，大多数低级内核（如 resize_image() 或 crop_image()）应该能够正确编译，没有图中断且支持动态形状。我们仍在解决剩余的边缘情况，朝着完整的函数支持和类迈进，您应该期待在下个版本中在该方面取得更多进展。