• 【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 ) 2^{(n-15+1)} 2(n−15+1)，具体的细节这里不说明。需要看时再百度。
    float类型在内存中的表示
     
    单独使用FP16：

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

      • 1
      • 2
      • 3
      • 4
      • 5
      • 6
      • 7
      • 8

    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()) 

      • 1
      • 2
      • 3

    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 的模型参数
         

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

    # 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()

    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17
    • 18
    • 19

    具体介绍如下。

    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。
After exiting autocast, calls f_float16.float() to use with d_float32 
g_float32 = torch.mm(d_float32, f_float16.float())

    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17

     
    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)

    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11

    3.3 autocast 作为装饰器

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

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

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

    • 1
    • 2
    • 3

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

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

    • 1
    • 2
    • 3
    • 4

     
    为了在其他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)

    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17
    • 18

    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_bactor，或者到了该增长的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

    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15

    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()

    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17
    • 18

     
    Gradient accumulation # 具体代码有待yolov5确认

    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()

    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17
    • 18
    • 19
    • 20
    • 21

     
    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()

    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17
    • 18
    • 19
    • 20
    • 21
    • 22
    • 23
    • 24
    • 25
    • 26
    • 27
    • 28
    • 29

    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()

    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17
    • 18
    • 19
    • 20
    • 21
    • 22
    • 23
    • 24

    5 混合精度使用注意事项

    • 尽量在 具有 Tensor Core 架构的 GPU 使用amp。
      在没有Tensor Core 架构的GPU 上使用amp，显存会明显减小，但速度会下降较多。具体的，在 Turing架构的 GTX 1660 上使用amp，运算时间增加了 一倍，显存不到原来的一半
    • 常数范围：为了保证计算不溢出，首先保证人工设定的常数不溢出。如 epsilon、INF 等
    • Dimension 最好是8的倍数：维度是8的倍数，性能最好
  • 相关阅读:
    主成分分析（机器学习）
    Docker下常规软件安装
    SparkCore
    MySQL--数据库的操作
    MyBatis的关联映射
    LCD DRM驱动框架分析一
    总结数据结构常用树
    腾讯云web应用防火墙与DDos高防包结合应用！
    HL7入门收集
    基于JAVA南京传媒学院门户网计算机毕业设计源码+系统+mysql数据库+lw文档+部署
  • 原文地址：https://blog.csdn.net/magic_ll/article/details/124689395