【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。
【pytorch记录】自动混合精度训练 torch.cuda.amp_第1张图片
其中,sign为表示正负,exponent位表示指数 2 ( n − 15 + 1 ) 2^{(n-15+1)} 2(n15+1),具体的细节这里不说明。需要看时再百度。
float类型在内存中的表示

单独使用FP16:

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

你可能感兴趣的:(pytorch,pytorch,python)