YOLOv5系列(三十二) 解读PyTorch的自动混合精度AMP(详尽)

官方文档：https://pytorch.org/docs/stable/notes/amp_examples.html

简单直观的理解

背景：

pytorch从1.6版本开始，已经内置了torch.cuda.amp，采用自动混合精度训练就不需要加载第三方NVIDIA的apex库了。本文主要从三个方面来介绍AMP：

一．什么是AMP?

二．为什么要使用AMP？

三．如何使用AMP?

四. 注意事项

## 一．什么是AMP?

默认情况下，大多数深度学习框架都采用32位浮点算法进行训练。２０１７年，NVIDIA研究了一种用于混合精度训练的方法，该方法在训练网络时将单精度（FP32）与半精度(FP16)结合在一起，并使用相同的超参数实现了与FP32几乎相同的精度。

在介绍AMP之前，先来理解下FP16与FP32，FP16也即半精度是一种计算机使用的二进制浮点数据类型，使用2字节存储。而FLOAT就是FP32。

其中，sign位表示正负，exponent位表示指数$2^{(}n-15+1(n=0))$，fraction位表示分数(m/1024)。

一般情况下，我们在pytorch中创建一个Tensor:

>>import torch
>>tensor1=torch.zeros(30,20)
>>tensor1.type()
'torch.FloatTensor'

>>tensor2=torch.Tensor([1,2])
>>tensor2.type()

'torch.FlatTensor'

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

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

默认Tensor是32bit floating point，这就是32位浮点型精度的tensor。

AMP(自动混合精度）的关键词有两个：自动，混合精度。

自动：Tensor的dtype类型会自动变化，框架按需自动调整tensor的dtype,当然有些地方还需手动干预。

混合精度：采用不止一种精度的Tensor，torch.FloatTensor和torch.HalfTensor

pytorch1.6的新包：torch.cuda.amp，是ＮVIDIA开发人员贡献到pytorch里的。只有支持tensor core的CUDA硬件才能享受到AMP带来的优势。Tensor core是一种矩阵乘累加的计算单元，每个tensor core时针执行64个浮点混合精度操作（FP16矩阵相乘和FP32累加）。

二、为什么要使用AMP?

​ 前面已介绍，AMP其实就是Float32与Float16的混合，那为什么不单独使用Float32或Float16，而是两种类型混合呢？原因是：在某些情况下Float32有优势，而在另外一些情况下Float16有优势。这里先介绍下FP16：

优势有三个：

１．减少显存占用；

２．加快训练和推断的计算，能带来多一倍速的体验；

３．张量核心的普及（NVIDIA　Tensor Core）,低精度计算是未来深度学习的一个重要趋势。

但凡事都有两面性，FP16也带来了些问题：１．溢出错误；２．舍入误差；

１．溢出错误：由于FP16的动态范围比FP32位的狭窄很多，因此，在计算过程中很容易出现上溢出和下溢出，溢出之后就会出现"NaN"的问题。在深度学习中，由于激活函数的梯度往往要比权重梯度小，更易出现下溢出的情况

2.舍入误差

舍入误差指的是当梯度过小时，小于当前区间内的最小间隔时，该次梯度更新可能会失败：

为了消除torch.HalfTensor也就是FP16的问题，需要使用以下两种方法：

１）混合精度训练

在内存中用FP16做储存和乘法从而加速计算，而用FP32做累加避免舍入误差。混合精度训练的策略有效地缓解了舍入误差的问题。

什么时候用torch.FloatTensor,什么时候用torch.HalfTensor呢？这是由pytorch框架决定的，在pytorch1.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

2)损失放大（Loss scaling)

即使了混合精度训练，还是存在无法收敛的情况，原因是激活梯度的值太小，造成了溢出。可以通过使用torch.cuda.amp.GradScaler，通过放大loss的值来防止梯度的underflow（只在BP时传递梯度信息使用，真正更新权重时还是要把放大的梯度再unscale回去）；

反向传播前，将损失变化手动增大$2^k$倍，因此反向传播时得到的中间变量（激活函数梯度）则不会溢出；

反向传播后，将权重梯度缩小$2^k$倍，恢复正常值。

三．如何使用AMP?

目前有两种版本：pytorch1.5之前使用的NVIDIA的三方包apex.amp和pytorch1.6自带的torch.cuda.amp

１．pytorch1.5之前的版本（包括１．５）

使用方法如下：

from apex import amp
model,optimizer = amp.initial(model,optimizer,opt_level="O1")   #注意是O,不是０
with amp.scale_loss(loss,optimizer) as scaled_loss:
    scaled_loss.backward()取代loss.backward()

其中,opt_level配置如下：

O0:纯FP32训练，可作为accuracy的baseline；

O1:混合精度训练（推荐使用），根据黑白名单自动决定使用FP16(GEMM,卷积）还是FP32（softmax)进行计算。

O2:几乎FP16，混合精度训练，不存在黑白名单　，除了bacthnorm，几乎都是用FP16计算；

O3:纯FP16训练，很不稳定，但是可以作为speed的baseline；

动态损失放大（dynamic loss scaling)部分，为了充分利用FP16的范围，缓解舍入误差，尽量使用最高的放大倍数$2^{2}4$,如果产生上溢出，则跳出参数更新，缩小放大倍数使其不溢出。在一定步数后再尝试使用大的scale来充分利用FP16的范围。

分布式训练：

import argparse
import apex import amp
import apex.parallel import convert_syncbn_model
import apex.parallel import DistributedDataParallel as DDP

定义超参数：
def parse():
    parser=argparse.ArgumentParser()
    parser.add_argument('--local_rank',type=int, default=0)  #local_rank指定了输出设备，默认为GPU可用列表中的第一个GPU，必须加上。
    ...
    args = parser.parser.parse_args()
    return args

主函数写：
def main():
    args = parse()
    torch.cuda.set_device(args.local_rank)  #必须写在下一句的前面
   torch.distributed.init_process_group(
       'nccl',
       init_method='env://')

导入数据接口，需要用DistributedSampler
    dataset = ...
    num_workers = 4 if cuda else 0
    train_sampler=torch.utils.data.distributed.DistributedSampler(dataset)
    loader = DataLoader(dataset, batch_size=batchsize, shuflle=False, num_worker=num_workers,pin_memory=cuda, drop_last=True, sampler=train_sampler)

定义模型：
net = XXXNet(using_amp=True)
net.train()
net= convert_syncbn_model(net)
device=torch.device('cuda:{}'.format(args.local_rank))
net=net.to(device)

定义优化器，损失函数，定义优化器一定要把模型搬运到GPU之上
apt = Adam([{'params':params_low_lr,'lr':4e-5},
    {'params':params_high_lr,'lr':1e-4}],weight_decay=settings.WEIGHT_DECAY)
crit = nn.BCELoss().to(device)

多GPU设置import torch.nn.parallel.DistributedDataParallel as DDP
net,opt = amp.initialize(net,opt,opt_level='o1')
net=DDP(net,delay_allreduce=True)loss使用方法：opt.zero_grad()with amp.scale_loss(loss, opt) as scaled_loss:    scaled_loss.backward()opt.step()加入主入口：if __name__ == '__main__':    main()无论是apex支持的DDP还是pytorch自身支持的DDP,都需使用torch.distributed.launch来使用，方法如下：CUDA_VISIBLE_DIVECES=1,2,4 python -m torch.distributed.launch --nproc_per_node=3 train.py1,2,4是GPU编号，nproc_per_node是指定用了哪些GPU,记得开头说的local_rank，是因为torch.distributed.launch会调用这个local_ran

分布式训练时保存模型注意点：

如果直接在代码中写torch.save来保存模型，则每个进程都会保存一次相同的模型，会存在写文件写到一半，会被个进程写覆盖的情况。如何避免呢?

可以用local_rank == 0来仅仅在第一个GPU上执行进程来保存模型文件。

虽然是多个进程，但每个进程上模型的参数值都是一样的，而默认代号为０的进程是主进程if arg.local_rank == 0:
    torch.save(xxx)

２．pytorch1.6及以上版本

有两个接口：autocast和Gradscaler

1. autocast

导入pytorch中模块torch.cuda.amp的类autocast

from torch.cuda.amp import autocast as autocast

model=Net().cuda()
optimizer=optim.SGD(model.parameters(),...)

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

  with autocast():
    output=model(input)
    loss = loss_fn(output,target)

  loss.backward()
  optimizer.step()

可以使用autocast的context managers语义（如上），也可以使用decorators语义。当进入autocast上下文后，在这之后的cuda ops会把tensor的数据类型转换为半精度浮点型，从而在不损失训练精度的情况下加快运算。而不需要手动调用.half(),框架会自动完成转换。

不过，autocast上下文只能包含网络的前向过程(包括loss的计算），不能包含反向传播，因为BP的op会使用和前向op相同的类型。

当然，有时在autocast中的代码会报错：

Traceback (most recent call last):
......
 File "/opt/conda/lib/python3.8/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 scaler type float but found c10:Half,应该是个bug,可在tensor上手动调用.float()来让type匹配。

２）GradScaler

使用前，需要在训练最开始前实例化一个GradScaler对象，例程如下：

from torch.cuda.amp import autocast as autocast

model=Net().cuda()
optimizer=optim.SGD(model.parameters(),...)

scaler = GradScaler() #训练前实例化一个GradScaler对象

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

    with autocast():　＃前后开启autocast
      output=model(input)
      loss = loss_fn(output,targt)

    scaler.scale(loss).backward()  #为了梯度放大
    #scaler.step()　首先把梯度值unscale回来，如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重，否则，忽略step调用，从而保证权重不更新。　　 scaler.step(optimizer)
    scaler.update()  #准备着，看是否要增大scaler

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

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

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

对于分布式训练，由于autocast是thread local的，要注意以下情形：

１）torch.nn.DataParallel：

以下代码分布式是不生效的

model = MyModel()
dp_model = nn.DataParallel(model)

with autocast():
    output=dp_model(input)
loss=loss_fn(output)

需使用autocast装饰model的forward函数

MyModel(nn.Module):
    @autocast()
    def forward(self, input):
        ...
        
#alternatively
MyModel(nn.Module):
    def forward(self, input):
        with autocast():
            ...


model = MyModel()
dp_model=nn.DataParallel(model)

with autocast():
    output=dp_model(input)
    loss = loss_fn(output)

2）torch.nn.DistributedDataParallel:

同样，对于多GPU,也需要autocast装饰model的forward方法，保证autocast在进程内部生效。

四. 注意事例：

在使用AMP时，由于报错信息并不明显，给调试带来了一定的难度。但只要注意以下一些点，相信会少走很多弯路。

１．判断GPU是否支持FP16，支持Tensor core的GPU（2080Ti,Titan,Tesla等），不支持的(Pascal系列）不建议；

1080Ti与2080Ti对比

gtx 1080ti:
半精度浮点数：0.17TFLOPS
单精度浮点数：11.34TFLOPS
双精度浮点数：0.33TFLOPS
rtx 2080ti:
半精度浮点数：20.14TFLOPS
单精度浮点数：10.07TFLOPS
双精度浮点数：0.31TFLOPS

半精度浮点数即FP16，单精度浮点数即FP32，双精度浮点数即FP64。
在不使用apex的pytorch训练过程中，一般默认均为单精度浮点数，从上面的数据可以看到1080ti和2080ti的单精度浮点数运算能力差不多，因此不使用apex时用1080ti和2080ti训练模型时间上差别很小。

使用apex时用1个2080ti训练时一个epoch是2h31min，两者时间几乎一样，但是却少用了一张2080ti。这是因为在pytorch训练中使用apex时，此时大多数运算均为半精度浮点数运算，而2080ti的半精度浮点数运算能力是其单精度浮点数运算能力的两倍

２．常数范围：为了保证计算不溢出，首先保证人工设定的常数不溢出。如epsilon,INF等；

３．Dimension最好是8的倍数:维度是８的倍数，性能最好；

４．涉及sum的操作要小心，容易溢出，softmax操作，建议用官方API，并定义成layer写在模型初始化里；

５．模型书写要规范：自定义的Layer写在模型初始化函数里，graph计算写在forward里；

６．一些不常用的函数，使用前要注册：amp.register_float_function(torch,‘sogmoid’)

７．某些函数不支持FP16加速，建议不要用；

８．需要操作梯度的模块必须在optimizer的step里，不然AMP不能判断grad是否为NaN

更加详细的解读

背景

自动混合精度。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）。

最后

你也可以使用我们提供的PyTorch项目规范来简化开发：

继承自DeepvacTrain类，在deepvac_config中设置config.amp = True即可。

在Gemfield的一个Conv2d和全连接层占主导的网络中，当开启AMP后，训练时显存的占用从11GB下降到了8GB。
官方解读：

https://zhuanlan.zhihu.com/p/348554267