【pytorch记录】自动混合精度训练 torch.cuda.amp

Nvidia 在Volta 架构中引入 Tensor Core 单元，来支持 FP32 和 FP16 混合精度计算。tensor core是一种矩阵乘累加的计算单元，每个tensor core 时执行64个浮点混合精度操作（FP16矩阵相乘和FP32累加）。
同年提出了一个pytorch 扩展apex，来支持模型参数自动混合精度训练。
Pytorch1.6版本以后，开始原生支持amp，即 torch.cuda.amp,是 nvidia开发人员贡献到pytorch里的，只有支持 tensor core 的 CUDA 硬件才能享受到 amp 带来的优势。

1 FP16半精度

FP16 和 FP32，是计算机使用的二进制浮点数据类型。
FP16 即半精度，使用2个字节。FP32 即Float。

其中，sign为表示正负，exponent位表示指数 $2^{(n-15+1)}$，具体的细节这里不说明。需要看时再百度。
float类型在内存中的表示
 
单独使用FP16：

• 优势：
  • 减小显存的占用，从而能支持更多的batchsize、更大模型和更大的输入尺寸 进行训练，有时反而带来精度上的提升
    加快训练和推理的计算，能加快一倍的速度
• 缺点：
  • 溢出问题：
    由于FP16的动态范围比FP32 的数值范围小很多，所以在计算过程中很容易出现上溢出和下溢出，然后就出现了"NAN"的问题。在深度学习中，由于激活函数的梯度往往比权重梯度小，更易出现下溢出问题。
    当第L层的梯度下溢出时，第L-1层已经以前的所有层的权重都无法更新
  • 舍入误差
    指 当梯度过小时，小于当前区间内的最小间隔时，该次梯度更新可能失败。比如
    FP16的 $2^{-3}+2^{-14}=2^{-3}$，此时就发生了舍入误差：在 $[2^{-3},2^{-2}]$ 间，比$2^{-3}$大的下一个数为 ($2^{-3}+2^{-13}$)
     

  import numpy as np
  
  a = np.array(2**(-3),dtype=np.float16)
  b = np.array(2**(-14),dtype=np.float16)
  c = a+b
  print(a)                # 0.125
  print('%f'%b)   # 0.000061
  print(c)                # 0.125
  

---

pytorch中的数据类型：

• 在pytorch中，一共有10中类型的tensor：
  torch.FloatTensor – 32bit floating point （pytorch默认创建的tensor的类型）
  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

  import torch 
   tensor = torch.zeros(20,20)
   print(tensor.type()) 
  

2 混合精度训练机制

自动混合精度（Automatic Mixed Precision, AMP）训练，是在训练一个数值精度为32的模型时，一部分算子的操作 数值精度为FP16，其余算子的操作精度为FP32。具体的哪些算子使用的精度，amp自动设置好了，不需要用户额外设置。
这样在不改变模型、不降低模型训练精度的前提下，可以缩短训练时间，降低存储需求，从而能支持更多的batchsize、更大模型和更大的输入尺寸 进行训练。
 
torch.cuda.amp给用户提供了很方便的混合精度训练机制，通过使用 amp.autocast 和 amp.GradScaler 来实现：

1. 用户不需要手动对模型参数的dtype，amp会自动为算子选择合适的数值精度
2. 在反向传播时，FP16的梯度数值溢出的问题，amp提供了梯度scaling操作，而且在优化器更新参数前，会自动对梯度 unscaling。所以对模型优化的超参数不会有任何的影响。

 
具体的实现流程如下：

正常的神经网络训练：前向计算loss、反向梯度计算、梯度更新。
混合精度训练：拷贝权重副本并转为FP16模型、前向计算loss、loss放大、反向梯度计算、梯度缩小、FP16 的梯度更新到FP32模型。

具体的amp的训练流程：

• 维护一个 FP32 数值精度模型的副本
• 在每个迭代
  • 拷贝并且转换成 FP16 模型。
  • 前向传播（FP16的模型参数）
    FP16的算子，直接计算操作；对 FP32 的算子，输入输出是FP16，计算的精度为FP32。反向时同理
  • loss 放大 s 倍
  • 反向传播，也就是反向梯度计算（FP16的模型参数和参数梯度）
  • 梯度乘以 1/s
  • 利用 FP16 的梯度更新 FP32 的模型参数
     

其中放大系数 s 的选择，选择一个常量是不合适的。因为loss和梯度的数值是变化的，所以 s 需要跟着 loss 来动态变化。
健康的loss 振荡中下降，因此 GradScaler 设计的 s 每隔 N 个 iteration 乘一个大于1的系数，在scale loss；

• 维护一个 FP32 数值精度模型的副本
• 在每个迭代
  1. 拷贝并且转换成 FP16 模型。
  2. 前向传播（FP16的模型参数）
  3. loss 放大 s 倍
  4. 反向传播，也就是反向梯度计算（FP16的模型参数和参数梯度）
  5. 检查是否有 inf 或者 nan 的参数梯度。如果有，降低 s，回到步骤1
  6. 梯度乘以 1/s
  7. 利用 FP16 的梯度更新 FP32 的模型参数
      

用户使用混合精度训练基本操作如下：

from torch.cuda.amp import GradScaler as GradScaler

# amp依赖Tensor core架构，所以model参数必须是cuda tensor类型
model = Net().cuda() optimizer = optim.SGD(model.parameters(), ...)
# GradScaler对象用来自动做梯度缩放 
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        # 在autocast enable 区域运行forward
        with autocast():
            # model做一个FP16的副本，forward
            output = model(input)
            loss = loss_fn(output, target)
            
        # 用scaler，scale loss(FP16)，backward得到scaled的梯度(FP16)
        scaler.scale(loss).backward()
        
        # scaler 更新参数，会先自动unscale梯度
        # 如果有nan或inf，自动跳过
        scaler.step(optimizer)
        
        # scaler factor更新
        scaler.update()

具体介绍如下。

3 aotucast

---

classs aotucast(device_type, enable=True, **kwargs)

• [device_type] (string) 表示是否使用 ‘cuda’ 或者 ‘cpu’ 设备
• [enabled] (bool，默认为True) 表示是否在区域中启用自动投射（自动转换）
• [dtype] (torch_dpython 类型) 表示使用 torch.float16/ torch.bfloat16
• [cache_enabled] (bool，默认为True) 表示是否使用 autocast 中的权重缓存

说明：

• autocast 的实例可以用作上下文管理器 或装饰器，设置区域以混合精度运行

---

3.1 autocast 算子

在pytorch中，在使用autocast的区域，会将部分算子自动转换成FP16 进行计算。只有CUDA算子有资格被自动转换。

• amp 自动转换成 FP16 的算子有：
  
• 自动转换成 FP32 的算子：
  
• 剩下没有列出的算子，像dot,add,cat…都是按数据中较大的数值精度，进行操作，即有 FP32 参与计算，就按 FP32，全是 FP16 参与计算，就是 FP16。

---

3.2 显示转换精度的情况

进入autocast-enabled 区域时，张量可以是任何类型。使用自动投射时，不应在模型或输入上调用 half() 或 bfloat16()。
但，作为上下文管理器使用时，混合精度计算enable区域得到的FP16数值精度的变量在enable区域外要显式的转换成FP32，否则使用过程中可能会导致类型不匹配的错误

# 在默认数据类型中创建一些张量（此处假定为FP32）
a_float32 = torch.rand((8, 8), device="cuda") 
b_float32 = torch.rand((8, 8), device="cuda") 
c_float32 = torch.rand((8, 8), device="cuda") 
d_float32 = torch.rand((8, 8), device="cuda")

with autocast():
    # torch.mm 是在 autocast算子的列表中，会转换为 FP16.
    # 输入为FP32, 但会以FP16精度运行计算，并输出FP16数据
    # 这个过程不需要手动设置
    e_float16 = torch.mm(a_float32, b_float32)
    # 也可以是混合输入类型
    f_float16 = torch.mm(d_float32, e_float16)

# 但 在退出 autocast 后，使用autocast区域生成的FP16变量时，就需要显示的转换成FP32。
g_float32 = torch.mm(d_float32, f_float16.float())

 
autocast也可以嵌套使用：

# 在默认数据类型中创建一些张量（此处假定为FP32）
a_float32 = torch.rand((8, 8), device="cuda") 
b_float32 = torch.rand((8, 8), device="cuda") 
c_float32 = torch.rand((8, 8), device="cuda") 
d_float32 = torch.rand((8, 8), device="cuda")

with autocast():
    e_float16 = torch.mm(a_float32, b_float32)
    with autocast(enabled=False):
    	f_float32 = torch.mm(c_float32, e_float16.float())
    g_float16 = torch.mm(d_float32, f_float32)

---

3.3 autocast 作为装饰器

这种情况一般使用分布式训练中。autocast 设计为 “thread local” 的，所以只在 main thread 上设 autocast 区域是不 work 的：

非分布式训练一般调用形式为：

model = MyModel() 
with autocast(): 	
	output = model(input)

 
分布式训练会用 nn.DataParalle() 或 nn.DistributedDataParallel，在创建model之后添加相应的代码，如下，但这样是不生效的，这里的autocast只在main thread中工作：

model = MyModel() 
DP_model = nn.DataParalle(model)  ## 添加
with autocast(): 	
	output = DP_model(input)

 
为了在其他thread上同时也生效，需要在定义网络结构中的 forward 也设置 autocast。有两种方式，添加装饰器、添加上下文管理器。

## 方式1：装饰器
class myModel(nn.Module):
@autocast()
	def forward(self, input):
		pass

## 方式2：上下文管理器
class myModule(nn.Module):
	def forward(self, input):
		with autocast():
			pass

## 主函数中调用
model = MyModel() 
DP_model = nn.DataParalle(model)  ## 添加

with autocast(): 	
	output = DP_model(input)

4 GradScaler 类

当使用了混合精度训练，存在无法收敛的情况，原因是激活梯度的值太小了，造成了溢出。可以通过使用 torch.cuda.amp.GradScaler，放大loss的值 来防止梯度的underflow。
torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)

• 【init_scale】 scale factor的初始值
• 【growth_factor】 每个scale factor 的增长系数
• 【backoff_factor】scale factor 下降系数
• 【growth_interval】每隔多个interval增长scale factor
• 【enabled】是否做scale

---

4.1 GradScaler 的方法

• scale(output) 方法
  对 outputs 成 scale factor，并返回。如果enabled=False，就直接返回
• step(optimizer, *args, **kwargs) 方法
  完成了两个功能：对梯度 unscale；检查梯度溢出，如果没有 nan/inf，就执行optimizer 的step，如果有就跳过
• update(new_scale=None) 方法
  update方法在每个 iteration 结束前都需要调用，如果参数更新跳过，会给scale factor 乘以 backoff_factor，或者到了该增长的iteration，就给scale factor 乘 growth_factor。也可以使用 new_scale 直接更新 scale factor.
   

例子：

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

---

4.2 GradScaler在梯度处理更多方面的应用

Gradient clipping

scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()

        # 先进行unscale 梯度，此时的clip threshold才能正确对梯度使用
        scaler.unscale_(optimizer)
        # clip梯度
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

        # unscale_() 已经被显式调用了，scaler执行step时就不再unscalse更新参数，有nan/inf也会跳过
        scaler.step(optimizer)
        scaler.update()

 
Gradient accumulation

scaler = GradScaler()

for epoch in epochs:
    for i, (input, target) in enumerate(data):
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
            # loss 根据 累加的次数归一一下
            loss = loss / iters_to_accumulate

        # scale 归一的loss 并backward  
        scaler.scale(loss).backward()

        if (i + 1) % iters_to_accumulate == 0:
            # may unscale_ here if desired 
            # (e.g., to allow clipping unscaled gradients)

            # step() and update() proceed as usual.
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

 
Gradient penalty

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
        # 防止溢出，在不是autocast 区域，先用scaled loss 得到 scaled 梯度
        scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss),
                                                 inputs=model.parameters(),
                                                 create_graph=True)
        # 梯度unscale
        inv_scale = 1./scaler.get_scale()
        grad_params = [p * inv_scale for p in scaled_grad_params]
        # 在autocast 区域，loss 加上梯度惩罚项
        with autocast():
            grad_norm = 0
            for grad in grad_params:
                grad_norm += grad.pow(2).sum()
            grad_norm = grad_norm.sqrt()
            loss = loss + grad_norm

        scaler.scale(loss).backward()

        # may unscale_ here if desired 
        # (e.g., to allow clipping unscaled gradients)

        # step() and update() proceed as usual.
        scaler.step(optimizer)
        scaler.update()

---

4.5 Multiple models

只需要使用一个scaler对多个模型操作，但scale(loss) 和 step(optimizer) 要分别执行
 

scaler = torc
h.cuda.amp.GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer0.zero_grad()
        optimizer1.zero_grad()
        with autocast():
            output0 = model0(input)
            output1 = model1(input)
            loss0 = loss_fn(2 * output0 + 3 * output1, target)
            loss1 = loss_fn(3 * output0 - 5 * output1, target)

        # 这里的retain_graph与amp无关，这里出现是因为在这个示例中，两个backward() 调用都共享图的一些部分。
        scaler.scale(loss0).backward(retain_graph=True)
        scaler.scale(loss1).backward()

        # 如果要检查或修改其拥有的参数的梯度，可以选择相应的优化器 进行显式取消缩放。
        scaler.unscale_(optimizer0)

        scaler.step(optimizer0)
        scaler.step(optimizer1)

        scaler.update()

5 混合精度使用注意事项

• 尽量在 具有 Tensor Core 架构的 GPU 使用amp。
  在没有Tensor Core 架构的GPU 上使用amp，显存会明显减小，但速度会下降较多。具体的，在 Turing架构的 GTX 1660 上使用amp，运算时间增加了 一倍，显存不到原来的一半
• 常数范围：为了保证计算不溢出，首先保证人工设定的常数不溢出。如 epsilon、INF 等
• Dimension 最好是8的倍数：维度是8的倍数，性能最好