PyTorch的自动混合精度（AMP）

PyTorch的自动混合精度（AMP）

转自：PyTorch的自动混合精度（AMP）

背景

PyTorch 1.6版本今天发布了，带来的最大更新就是自动混合精度。release说明的标题是：

1. Stable release of automatic mixed precision (AMP).
2. New Beta features include a TensorPipe backend for RPC, memory profiler,
3. and several improvements to distributed training for both RPC and DDP.

可见自动混合精度正是 PyTorch 1.6 的最大更新。这就带来了几个问题：

1. 什么是自动混合精度训练？
2. 为什么需要自动混合精度？
3. 如何在 PyTorch 中使用自动混合精度？

什么是自动混合精度训练？

我们知道神经网络框架的计算核心是 Tensor，也就是那个从 scaler -> array -> matrix -> tensor 维度一路丰富过来的 tensor。在 PyTorch 中，我们可以这样创建一个 Tensor：

>>> import torch

>>> gemfield = torch.zeros(70,30)
>>> gemfield.type()
'torch.FloatTensor'

>>> syszux = torch.Tensor([1,2])
>>> syszux.type()
'torch.FloatTensor'

可以看到默认创建的 tensor 都是 FloatTensor 类型。而在 PyTorch 中，一共有 10 种类型的 tensor：

• torch.FloatTensor (32-bit floating point)
• torch.DoubleTensor (64-bit floating point)
• torch.HalfTensor (16-bit floating point 1)
• torch.BFloat16Tensor (16-bit floating point 2)
• torch.ByteTensor (8-bit integer (unsigned))
• torch.CharTensor (8-bit integer (signed))
• torch.ShortTensor (16-bit integer (signed))
• torch.IntTensor (32-bit integer (signed))
• torch.LongTensor (64-bit integer (signed))
• torch.BoolTensor (Boolean)

由此可见，默认的 Tensor 是 32-bit floating point，这就是 32 位浮点型精度的 Tensor。

自动混合精度的关键词有两个：自动、混合精度。这是由 PyTorch 1.6 的torch.cuda.amp 模块带来的：

from torch.cuda.amp import autocast as autocast

混合精度预示着有不止一种精度的 Tensor，那在 PyTorch 的 AMP 模块里是几种呢？2种：torch.FloatTensor 和 torch.HalfTensor；

自动预示着 Tensor 的 dtype 类型会自动变化，也就是框架按需自动调整 tensor 的 dtype（其实不是完全自动，有些地方还是需要手工干预）；

torch.cuda.amp 的名字意味着这个功能只能在 cuda 上使用，事实上，这个功能正是 NVIDIA 的开发人员贡献到 PyTorch 项目中的。而只有支持 Tensor core 的 CUDA 硬件才能享受到 AMP 的好处（比如2080ti显卡）。Tensor Core 是一种矩阵乘累加的计算单元，每个 Tensor Core 每个时钟执行 64 个浮点混合精度操作（FP16 矩阵相乘和 FP32 累加），英伟达宣称使用 Tensor Core 进行矩阵运算可以轻易的提速，同时降低一半的显存访问和存储。

因此，在 PyTorch 中，当我们提到自动混合精度训练，我们说的就是在 NVIDIA 的支持 Tensor core 的 CUDA 设备上使用 torch.cuda.amp.autocast （以及torch.cuda.amp.GradScaler）来进行训练。咦？为什么还要有 torch.cuda.amp.GradScaler？

为什么需要自动混合精度？

这个问题其实暗含着这样的意思：为什么需要自动混合精度，也就是 torch.FloatTensor 和 torch.HalfTensor 的混合，而不全是 torch.FloatTensor？或者全是 torch.HalfTensor？

如果非要以这种方式问，那么答案只能是，在某些上下文中 torch.FloatTensor 有优势，在某些上下文中 torch.HalfTensor 有优势呗。答案进一步可以转化为，相比于之前的默认的 torch.FloatTensor，torch.HalfTensor 有时具有优势，有时劣势不可忽视。

torch.HalfTensor 的优势就是存储小、计算快、更好的利用 CUDA 设备的 Tensor Core。因此训练的时候可以减少显存的占用（可以增加 batchsize 了），同时训练速度更快；

torch.HalfTensor 的劣势就是：数值范围小（更容易 Overflow / Underflow）、舍入误差（Rounding Error，导致一些微小的梯度信息达不到 16bit 精度的最低分辨率，从而丢失）。

可见，当有优势的时候就用 torch.HalfTensor，而为了消除 torch.HalfTensor 的劣势，我们带来了两种解决方案：

1，梯度 scale，这正是上一小节中提到的 torch.cuda.amp.GradScaler，通过放大 loss 的值来防止梯度的 underflow（这只是BP的时候传递梯度信息使用，真正更新权重的时候还是要把放大的梯度再unscale回去）；

2，回落到 torch.FloatTensor，这就是混合一词的由来。那怎么知道什么时候用 torch.FloatTensor，什么时候用半精度浮点型呢？这是 PyTorch 框架决定的，在PyTorch 1.6 的 AMP 上下文中，如下操作中 tensor 会被自动转化为半精度浮点型的 torch.HalfTensor：

1. __matmul__
2. addbmm
3. addmm
4. addmv
5. addr
6. baddbmm
7. bmm
8. chain_matmul
9. conv1d
10. conv2d
11. conv3d
12. conv_transpose1d
13. conv_transpose2d
14. conv_transpose3d
15. linear
16. matmul
17. mm
18. mv
19. prelu

如何在 PyTorch 中使用自动混合精度

答案就是 autocast + GradScaler。

1，autocast

正如前文所说，需要使用 torch.cuda.amp 模块中的 autocast 类。使用也是非常简单的：

from torch.cuda.amp import autocast as autocast

# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

for input, target in data:
    optimizer.zero_grad()

    # 前向过程(model + loss)开启 autocast
    with autocast():
        output = model(input)
        loss = loss_fn(output, target)

    # 反向传播在autocast上下文之外
    loss.backward()
    optimizer.step()

可以使用 autocast 的 context managers 语义（如上所示），也可以使用 decorators 语义。 当进入 autocast 的上下文后，上面列出来的那些 CUDA ops 会把tensor 的 dtype 转换为半精度浮点型，从而在不损失训练精度的情况下加快运算。刚进入 autocast 的上下文时，tensor 可以是任何类型，你不要在 model 或者input 上手工调用 .half() ，框架会自动做，这也是自动混合精度中“自动”一词的由来。

另外一点就是，autocast 上下文应该只包含网络的前向过程（包括 loss 的计算），而不要包含反向传播，因为 BP 的 op 会使用和前向 op 相同的类型。

还有的时候呀，你的代码在 autocast 上下文中会报如下的错误：

Traceback (most recent call last):
......
  File "/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
    result = self.forward(*input, **kwargs)
......
RuntimeError: expected scalar type float but found c10::Half

对于 RuntimeError: expected scalar type float but found c10::Half，这估计是个 bug。你可以在 tensor 上手工调用 .float() 来让 type 匹配。

2，GradScaler

但是别忘了前面提到的梯度 scaler 模块呀，需要在训练最开始之前实例化一个 GradScaler 对象。因此 PyTorch 中经典的 AMP 使用方式如下：

from torch.cuda.amp import autocast as autocast

# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# 在训练最开始之前实例化一个GradScaler对象
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # 前向过程(model + loss)开启 autocast
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # Scales loss. 为了梯度放大.
        scaler.scale(loss).backward()

        # scaler.step() 首先把梯度的值unscale回来.
        # 如果梯度的值不是 infs 或者 NaNs, 那么调用optimizer.step()来更新权重,
        # 否则，忽略step调用，从而保证权重不更新（不被破坏）
        scaler.step(optimizer)

        # 准备着，看是否要增大scaler
        scaler.update()

scaler 的大小在每次迭代中动态的估计，为了尽可能的减少梯度 underflow，scaler 应该更大；但是如果太大的话，半精度浮点型的 tensor 又容易 overflow（变成 inf 或者 NaN）。所以动态估计的原理就是在不出现 inf 或者 NaN 梯度值的情况下尽可能的增大 scaler 的值——在每次 scaler.step(optimizer) 中，都会检查是否又 inf 或 NaN 的梯度出现：

1. 如果出现了 inf 或者 NaN，scaler.step(optimizer) 会忽略此次的权重更新（optimizer.step() )，并且将 scaler 的大小缩小（乘上 backoff_factor ）；

2. 如果没有出现 inf 或者 NaN，那么权重正常更新，并且当连续多次（ growth_interval 指定）没有出现 inf 或者 NaN，则 scaler.update() 会将 scaler 的大小增加（乘上 growth_factor ）。