翻译自 PyTorch 官方 Blog 并添加少量修改 MMCV 自 v1.4.7 开始提供 PyTorch1.11 预编译包
前言
近日，PyTorch 推出 PyTorch 1.11，本次亮点可总结为如下 ：
- TorchData：一个新的用于通用模块化数据加载的库，可以轻松构建灵活、高性能的数据 pipeline
- functorch：一个类 JAX 的向 PyTorch 添加可组合函数转换的库
- DDP 静态图优化正式可用
TorchData
TorchData 是一个通用模块化数据加载库，用于轻松构建灵活且高性能的数据 pipeline。根据社区反馈，PyTorch 发现现有的 DataLoader 将太多的功能耦合在一起，难以扩展；此外，不同的用例通常必须重写相同的数据加载程序，造成代码冗余。为了解决这些问题，PyTorch 设计了 TorchData，其目标是通过 Iterable-style 和 Map-style 的 DataPipe 来实现可组合的数据加载。DataPipe 用来替代之前版本的 Dataset，并内置了大量数据相关操作，如打开文件、解析文本、转换样本、缓存、混洗和批处理等。
DataPipe 接受 Python 的一些访问函数，例如 __iter__ 和 __getitem__，前者用于 IterDataPipe，后者用于 MapDataPipe，它们会返回一个新的访问函数。您可以将多个 DataPipe 连在一起，形成数据 pipeline，以执行必要的数据转换工作。PyTorch 目前提供了 50+ DataPipe。
在 PyTorch1.11 中，一些 PyTorch 库已经将数据集迁移到 DataPipe。在 TorchText 中提供的数据集是使用 DataPipe 实现的，其 SST-2 二进制文本分类教程的一部分 演示了如何使用 DataPipe 预处理数据。在 TorchVision 和 TorchRec 中还有其他数据集的原型实现，您可以参考官方示例。
TorchData 文档 已经上线，它包含一个教程，教程介绍了如何使用 DataPipe、将它们与 DataLoader 一起使用以及如何实现自定义 Datapipe。
下面我们举一个例子来了解一下 DataPipe：
from torchdata.datapipes.iter import IterDataPipe
import torchdata.datapipes as dp
@functional_datapipe("map")
class MapperIterDataPipe(IterDataPipe):
def __init__(self, source_dp: IterDataPipe, fn) -> None:
super().__init__()
self.dp = source_dp
self.fn = fn
def __iter__(self):
for d in self.dp:
yield self.fn(d)
# Using functional form (recommended)
datapipes1 = dp.iter.FileOpener(['a.file', 'b.file']).map(fn=decoder).这个 MapperIterDataPipe 看起来有点像 Iterable-style dataset，因为它也实现了 __iter__ 方法，但它并不是描述数据本身，而是描述对数据的变换。可以看到，在 __init__ 里，它接收两个参数，一个是源 DataPipe，一个是对数据进行转换的函数，然后在 __iter__ 里，它对 DataPipe 的每一个元素依次做转换，并用 yield 返回，这种 lazy 模式使得整体的 DataPipe 是内存友好的。另外，MapperIterDataPipe 可以通过一个装饰器 functional_datapipe 注册为函数（map），这样就可以通过 DataPipe.map(...) 的形式使用该 DataPipe。
functorch
PyTorch 官方宣布推出 functorch 的首个 beta 版本，该库受到 Google JAX 的极大启发。functorch 是一个向 PyTorch 添加可组合函数转换的库，旨在提供可组合的 vmap（向量化）和 autodiff 转换，可与 PyTorch 模块和 PyTorch autograd 一起使用，并具有良好的动态模式（eager-mode）性能。
可组合的函数转换可以帮助解决当前在 PyTorch 中难以实现的许多用例：
- 计算每个样本的梯度
- 单机运行多个模型的集成
- 在元学习（MAML）内循环中高效地批处理任务
- 高效地计算雅可比矩阵（Jacobians）和海森矩阵（Hessians）
vmap（向量化）、vjp（反向模式 autodiff）和 jvp（前向模式 autodiff）转换的组合使得用户毫不费劲地表达上述内容，无需为每个转换设计单独的库。
例如计算雅各比矩阵：
import torch
import torch.nn as nn
import torch.nn.functional as F
from functools import partial
from functorch import vmap, vjp
D = 16
weight = torch.randn(D, D)
bias = torch.randn(D)
x = torch.randn(D)
unit_vectors = torch.eye(D)
def predict(weight, bias, x):
return F.linear(x, weight, bias).tanh()
# PyTorch 方式
def compute_jac(xp):
jacobian_rows = [torch.autograd.grad(predict(weight, bias, xp), xp, vec)[0]
for vec in unit_vectors]
return torch.stack(jacobian_rows)
jacobian = compute_jac(x)
# functorch 方式
_, vjp_fn = vjp(partial(predict, weight, bias), x)
ft_jacobian, = vmap(vjp_fn)(unit_vectors)可以看到 functorch 方式用 vmap 替代了 for 循环，而 vmap 是经过优化的并行计算，因而可以极大地提高运行速度，同时 functorch 没有使用 autograd.grad 而是使用了 vjp，这是因为 vmap 不能直接用于 autograd.grad。另外，在 output size 大于 input size 时，还可以使用 jvp 替代 vjp 以进一步提高效率。
DDP 静态图
DDP 静态图假设用户的模型在每次迭代中都使用相同的一组已使用或未使用的参数，因此它对一些相关状态的了解是确定的，例如哪些 hook 将被触发、触发的次数以及第一次迭代后的梯度计算就绪顺序。
静态图在第一次迭代中缓存这些状态，因此它可以支持 DDP 在以往版本中无法支持的功能，例如无论是否有未使用的参数，在相同参数上支持多个激活检查点。当存在未使用的参数时，静态图功能也会应用性能优化，例如避免遍历图在每次迭代中搜索未使用的参数，并启用动态分桶（bucketing）顺序。DDP 静态图中的这些优化为一些推荐模型带来了 10% 的 QPS 增益。
要启用静态图，只需在 DDP 中设置 static_graph=True，如下所示：
ddp_model = DistributedDataParallel(model, static_graph=True)MMCV 支持
MMCV 自 v1.4.7 开始提供 PyTorch1.11 预编译包，安装命令如下：
pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu115/torch1.11/index.html
pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.11/index.html
# 下面这条命令只在 Linux 平台有效，因为 torch 1.11.0 没有提供 cuda 10.2 在 Windows 平台 的预编译包，所以 mmcv 也没有提供
pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu102/torch1.11/index.htmlMMCV GitHub 地址：https://github.com/open-mmlab/mmcv
欢迎大家访问 MMCV GitHub 主页体验~
https://github.com/open-mmlab/mmcvgithub.com/open-mmlab/mmcv
