pytorch混精度训练AutoCast与GradScaler

一、

autocast是pytorch实现的一种用于降低训练时显存消耗的技术。（仅在GPU上训练时可使用）

它的原理是用更短的总位数来保存浮点数，能够有效将显存消耗降低，从而设置更大的batch来加速训练。

但这样会导致有效位数减少，不可避免地造成精度的丢失，最终模型的收敛效果也会变差。

因此，使用混精度训练后通常还需要进行后续的更精细的训练让模型收敛效果更好。

 

二、

autocast通常与GradScaler一起使用

因为autocast会损失部分精度，从而导致梯度消失的问题，并且经过中间层时可能计算得到inf导致最终loss出现nan。

所以我们通常将GradScaler与autocast配合使用来对梯度值进行一些放缩，来缓解上述的一些问题。

示例：

from torch.cuda.amp import autocast, GradScaler

dataloader = ...
model = Model.cuda(0)
optimizer = ...
scheduler = ...
scaler = GradScaler()  # 新建GradScale对象，用于放缩
for epoch_idx in range(epochs):
    for batch_idx, (dataset) in enumerate(dataloader):
        optimizer.zero_grad()
        dataset = dataset.cuda(0)
        with autocast():  # 自动混精度
            logits = model(dataset)
            loss = ...
        scaler.scale(loss).backward()  # scaler实现的反向误差传播
        scaler.step(optimizer)  # 优化器中的值也需要放缩
        scaler.update()  # 更新scaler
　　scheduler.step()

...

 

三、

使用autocast技术进行混精度训练时loss经常会出现'nan'

有以下三种成因：

①精度损失，有效位数减少，导致输出时数据末位的值被省去，最终出现nan的现象。该情况可以使用GradScaler（上文所示）来解决。

②损失函数中使用了log等形式的函数，或是变量出现在了分母中，并且训练时，该数值变得非常小时，混精度可能会让该值更接近0或是等于0，导致了数学上的log(0)或是x/0的情况出现，从而出现'inf'或'nan'的问题。

这种时候需要针对该问题设置一个确定值。例如：当log(x)出现-inf的时候，我们直接将输出中该位置的-inf设置为-100，即可解决这一问题。

③模型内部存在的问题，比如模型过深，本身梯度回传时值已经非常小。这种问题难以解决。