欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

【PyTorch】torch.cuda.amp自动混合精度训练

程序员文章站 2024-02-28 19:38:34
...

一、什么是amp?

amp:Automatic mixed precision,自动混合精度,可以在神经网络推理过程中,针对不同的层,采用不同的数据精度进行计算,从而实现节省显存和加快速度的目的。

自动混合精度的关键词有两个:自动、混合精度。这是由PyTorch 1.6的torch.cuda.amp模块带来的:

from torch.cuda import amp

混合精度预示着有不止一种精度的Tensor,那在PyTorch的AMP模块里是几种呢?2种:torch.FloatTensor(浮点型 32位)和torch.HalfTensor(半精度浮点型 16位);

自动预示着Tensor的dtype类型会自动变化,也就是框架按需自动调整tensor的dtype(其实不是完全自动,有些地方还是需要手工干预);

注意:

  1. torch.cuda.amp 的名字意味着这个功能只能在cuda上使用。
  2. torch默认的tensor精度类型是torch.FloatTensor

二、为什么需要自动混合精度(amp)?

也可以这么问:为什么需要自动混合精度,也就是torch.FloatTensor和torch.HalfTensor的混合,而不全是torch.FloatTensor?或者全是torch.HalfTensor?

原因: 在某些上下文中torch.FloatTensor有优势,在某些上下文中torch.HalfTensor有优势。

torch.HalfTensor:

  1. torch.HalfTensor的优势就是存储小、计算快、更好的利用CUDA设备的Tensor Core。因此训练的时候可以减少显存的占用(可以增加batchsize了),同时训练速度更快;
  2. 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框架决定的,AMP上下文中,一些常用的操作中tensor会被自动转化为半精度浮点型的torch.HalfTensor(如:conv1d、conv2d、conv3d、linear、prelu等)

三、如何在PyTorch中使用自动混合精度?

答案是 autocast + GradScaler

1、autocast

使用torch.cuda.amp模块中的autocast 类。

from torch.cuda import amp

# 创建model,默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 判断能否使用自动混合精度
enable_amp = True if "cuda" in device.type else False

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

    # 前向过程(model + loss)开启 autocast
    with amp.autocast(enabled=enable_amp):
        output = model(input)
        loss = loss_fn(output, target)

    # 反向传播在autocast上下文之外
    loss.backward()
    optimizer.step()

注意:

  1. 当进入autocast,自动将torch.FloatTensor类型转化为torch.HalfTensor,而不需要手动设置model.half()/input.half,框架会自动做,这也是自动混合精度中“自动”一词的由来。
  2. autocast上下文应该只包含网络的前向过程(包括loss的计算),而不要包含反向传播。

2、GradScaler

这里GradScaler就是第二小节中提到的梯度scaler模块,需要在训练最开始之前使用amp.GradScaler实例化一个GradScaler对象。

from torch.cuda import amp

# 创建model,默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 判断能否使用自动混合精度
enable_amp = True if "cuda" in device.type else False
# 在训练最开始之前实例化一个GradScaler对象
scaler = amp.GradScaler(enabled=enable_amp)

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

        # 前向过程(model + loss)开启 autocast
        with amp.autocast(enabled=enable_amp):
            output = model(input)
            loss = loss_fn(output, target)
            
        # 1、Scales loss.  先将梯度放大 防止梯度消失
        scaler.scale(loss).backward()

        # 2、scaler.step()   再把梯度的值unscale回来.
        # 如果梯度的值不是 infs 或者 NaNs, 那么调用optimizer.step()来更新权重,
        # 否则,忽略step调用,从而保证权重不更新(不被破坏)
        scaler.step(optimizer)

        # 3、准备着,看是否要增大scaler
        scaler.update()
        
        # 正常更新权重
        optimizer.zero_grad()

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

注意:
再强调一点,amp只能在GPU环境下使用,因为一来amp是写在torch.cuda中的函数,而且amp的中的 amp.GradScaler和amp.autocast函数构造是这样的:
amp.GradScaler:


    def __init__(self,
                 init_scale=2.**16,
                 growth_factor=2.0,
                 backoff_factor=0.5,
                 growth_interval=2000,
                 enabled=True):
        if enabled and not torch.cuda.is_available():
            warnings.warn("torch.cuda.amp.GradScaler is enabled, but CUDA is not available.  Disabling.")
            self._enabled = False
        else:
            self._enabled = enabled

amp.autocast:

 def __init__(self, enabled=True):
        if enabled and not torch.cuda.is_available():
            warnings.warn("torch.cuda.amp.autocast only affects CUDA ops, but CUDA is not available.  Disabling.")
            self._enabled = False
        else:
            self._enabled = enabled

可以看出,如果不是GPU环境的话,这两个函数是不会正常执行的。

四、多GPU训练

单卡训练的话上面的代码已经够了。要是想多卡跑的话仅仅这样还不够,会发现在forward里面的每个结果都还是float32的,怎么办?

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, input_data_c1):
    	with autocast():
    		# code
    	return

只要把model中的forward里面的代码用autocast代码块方式运行就好了。

相关标签: PyTorch pytorch