https://zhuanlan.zhihu.com/p/165152789
PyTorch 1.6版本今天发布了，带来的最大更新就是自动混合精度。release说明的标题是：
- Stable release of automatic mixed precision (AMP).
- New Beta features include a TensorPipe backend for RPC, memory profiler,
- and several improvements to distributed training for both RPC and DDP.
可见自动混合精度正是PyTorch 1.6的最大更新。这就带来了几个问题：
- 什么是自动混合精度训练？
- 为什么需要自动混合精度？
- 如何在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：
- __matmul__
- addbmm
- addmm
- addmv
- addr
- baddbmm
- bmm
- chain_matmul
- conv1d
- conv2d
- conv3d
- conv_transpose1d
- conv_transpose2d
- conv_transpose3d
- linear
- matmul
- mm
- mv
- 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）。
最后
你也可以使用我们提供的PyTorch项目规范来简化开发：
https://github.com/deepVAC/deepvac/github.com
继承自DeepvacTrain类，在deepvac_config中设置config.amp = True即可。
在Gemfield的一个Conv2d和全连接层占主导的网络中，当开启AMP后，训练时显存的占用从11GB下降到了8GB。而速度......