PyTorch1.13 亮点一览，TorchEval、MultiPy 、TorchSnapshot 新库大解读

前言

近日，PyTorch 发布了 PyTorch 1.13 版本，本次新版发布主要包含以下亮点：

• 稳定版本的 BetterTransformer API，无需修改模型即可支持通用 Transformer 模型在推理时候使用快速路径执行 （fastpath execution）。除此之外还对 Transfomer 常用尺寸的 Add + Matmul 线性代数核函数进行了加速，以及默认使用嵌套 Tensor。

• 引入新的 CUDA 版本：11.7，同时弃用旧的 CUDA 版本：10.2 和 11.3。

• functorch 转移到 PyTorch 库内，不再需要安装单独的 functorch 包。

• 支持 Apple M1 系列芯片的原生构建。

随着 PyTorch 1.13 的发布，MMCV 从 1.7.0 版本也开始为 PyTorch 1.13 提供预编译包，并且提供了 macOS 的预编译包，欢迎大家安装使用：

pip install -U openmim
mim install mmcv-full

MMCV GitHub 主页：https://github.com/open-mmlab/mmcv

PyTorch 库的新变化

在发布 PyTorch 1.13 版本的同时，还对 PyTorch 相关库进行了一系列改进，此次 PyTorch 库的更新包括：

• TorchAudio，新增 Hybrid Demucs 模型和 SUPERB 基准数据集，以及用于编码媒体文件的 StreamWriter 类。

• TorchData，进一步整合了 DataLoader2 的 API，扩展了对通过 DataPipes 从其他云存储提供商加载数据的支持，包括 AWS、谷歌云存储和 Azure。

• TorchEval，提供常见评测指标，为 PyTorch 引入原生的评测指标支持。（本周 OpenMMLab 也发布了评测库 MMEval，详情见**往期内容**）

• torch::deploy，提供了一种在单个进程中使用多个独立 Python 解释器运行的方法，而无需共享全局解释器锁。

• TorchMultimodal，一个用于大规模训练 SoTA 多任务和多模态模型的 PyTorch 算法库。

• TorchRec，简化了优化器融合和分布式切片相关的 API，对集合通信 Tensor 应用量化和混合精度以提高训练效率。

• TorchSnapshot，一个用于 PyTorch 的高性能、高内存效率的 checkpoint 库，设计时考虑到了大型、复杂的分布式工作负载。

• TorchVision，引入了一个新的模型注册 API，新增两个视频分类模型，增加了 Swin Transformer V2 模型架构和 MaxViT，以及新增了 Simple Copy-Paste 数据增强策略。

• Torch-TensorRT，新增 torch.fx 转换到 TensorRT 的前端。

• TorchX，增加了新的列表 API，实验跟踪，弹性训练和改进的调度器支持。

受限于篇幅，本文将只对新增的几个 PyTorch 库进行展开介绍：TorchEval，torch::deploy 和 TorchSnapshot。

TorchEval

TorchEval 是 PyTorch 新发布的一个评测库，具有以下特点：

• 包含了丰富的评测指标，通过 PyTorch 尽可能地利用向量化和 GPU 加速。

• 使用 torch.distributed 与分布式训练工具无缝集成。

• 设计时考虑到了可扩展性：可以自由地创建自己的评测指标。

• 为基于 PyTorch 的模型提供分析内存和计算要求的工具。

TorchEval 中的评测指标提供了两种使用方式，分别是基于类的接口和函数式的接口。函数式的接口可以在 torcheval.metrics.functional 中找到，在使用单个进程评测时使用函数式接口比较方便。当需要多进程评测时候，应该使用 statefull 的类接口，以支持多进程之间的状态同步。

目前 TorchEval 已经支持了一些基础常用的分类、回归和排序的评测指标，并且提供了单进程和多进程评测的示例。

torch::deploy （MultiPy）

torch::deploy（MultiPy）是一个 C++ 库，使你能够在生产环境中运行 Eager 模式的 PyTorch 模型，而无需对模型进行任何修改以支持导出可部署的静态图。

Python 已经成为训练深度神经网络广泛使用的语言，然而当模型被用于推理时，它们通常被从 Python 程序中提取为 TensorFlow Graph 或 TorchScript，以满足推理的性能要求以及方便打包。

但是这个提取过程在 PyTorch 中可能很耗时，这阻碍了快速的原型设计。 MultiPy 展示了如何能够在使用 Python 进行推理的同时，满足性能要求和打包限制。

MulitPy 定义了一种新的模型容器格式，其中包含本地 Python 代码和数据。这种方法不需要模型的静态图提取步骤，简化了模型部署的过程，并使之更容易整合现有的 Python 库。特别的是，MultiPy 提出了一种在单个进程中使用多个 Python 解释器的方法，摆脱 Python GIL 的限制以实现可扩展的推理。

MultiPy 在 GitHub 上的一些流行的 PyTorch 模型上进行了评估，展示了它们如何被打包成推理格式，并与 TorchScript 比较了它们的性能。对于较大的模型，MultiPy 打包的 Python 模型的性能与 TorchScript 相同，而对于较小的模型，多解释器方法确保推理仍然是可扩展的。

如下伪代码展示了如何使用 torch::deploy 调用 Python 模型进行推理：

torch::deploy::InterpreterManager manager(4);

// access one of the 4 interpreters
auto I = manager.acquireOne();

// run infer from your_model.py
I.global("your_model", "infer")({at::randn({10, 240, 320})});

TorchSnapshot

TorchSanpshot 是一个为 PyTorch 应用专门开发的 checkpoint 工具，设计时考虑到了大型、复杂的分布式工作负载。具有以下特点：

1. 高性能

• TorchSnapshot 提供了一个采用各种优化加速的 checkpoint 实现，包括大多数张量类型的零拷贝序列化、设备到主机拷贝和存储 I/O 的 Overlap、并行化存储 I/O。

• TorchSnapshot 通过在所有 rank 上分配写负荷，大大加快了分布式数据并行保存 checkpoint 的速度。

• 当主机内存充足时，TorchSnapshot 允许在所有存储 I/O 完成之前恢复训练，减少保存 checkpoint 所阻断的时间。

2. 高效的内存使用

• TorchSnapshot 的内存使用情况适应于主机的可用资源，大大减少了保存和加载 checkpoint 时出现内存不足问题的情况。

• TorchSnapshot 支持对快照中的单个对象进行有效的随机访问，即使快照被存储在云对象存储中。

3. 高可用性

• 简单的 API，在分布式和非分布式工作负载之间是一致的。

• 与常用的云对象存储系统开箱即用的集成，比如 S3 和谷歌云存储。

• 对于支持的工作负载，在并行进程数量变化时自动重新分区。

4. 安全性

• 安全的张量序列化，不依赖 pickle （正在支持中）。

结语

此次 PyTorch 1.13 的更新，除了 PyTorch 本身 API 的变化以及新 Feature 的增加以外，也有许多 PyTorch 库也进行了更新。这些 PyTorch 库涉及领域从视觉、语音到推荐算法和多模态，从评测、数据工具到 checkpoint 和部署工具。这些 PyTorch 库的更新表明，PyTorch 正专注于在许多领域开发通用和可扩展的 API，使社区更容易在 PyTorch 上构建生态系统项目。

本次提到的更新中，关于部署工具的更新有两个，分别是 Torch-TensorRT 新增 torch.fx 前端，和 torch::deploy（MultiPy）Beta 版本的发布。

PyTorch 由于其动态图的特性，在模型部署方面天然的弱于拥有静态图特性的深度学习框架。PyTorch 一直在尝试寻找好的解决方法，目前可以看到在动态图静态化方面，已经有 TorchScript 和 torch.fx 两个方案。而 MultiPy 或许是提出了另外一个思路，跳过动态图静态化这一步，直接将 Python 程序部署推理，期待 MultiPy 后续的发展。

另外本次 TorchEval 的发布也是一大亮点，截止目前，今年已经有三个开源的算法评测库发布，按照时间顺序分别是 huggingface/evaluate，pytorch/torcheval 和 open-mmlab/mmeval，再加上更早时候的 Lightning-AI/metrics，目前已有四个专注于算法评测的工具库。随着模型训练工具链逐步完善，模型评测工具链的价值也被大家所重视。

在这些开源评测库之中，OpenMMLab 开源的 MMEval 定位为统一开放的跨框架算法评测库。与 TorchEval 相比，MMEval 具有多框架支持的能力，并且提供了更为丰富和细分的评测指标，具体可以参考往期内容。MMEval 目前已支持 20+ 细分任务领域的评测指标，更多的评测指标正在火热添加中。同时我们也发布了一些为 MMEval 添加 Metric 的任务，欢迎大家参与：https://github.com/open-mmlab/mmeval/issues/50

参考链接

• https://pytorch.org/blog/PyTorch-1.13-release/

• https://pytorch.org/blog/new-library-updates-in-pytorch-1.13/

• https://github.com/pytorch/torcheval

• https://github.com/pytorch/multipy

• https://arxiv.org/pdf/2104.00254.pdf

• https://github.com/pytorch/torchsnapshot