deep-learning - 使用 pytorch 进行混合精度训练是否需要 GradScaler？

Question

所以去 AMP: Automatic Mixed Precision Training Normal 网络教程，我发现有两个版本，Automatic 和 GradScaler。我只想知道在训练中使用 GradScaler 是否可取/必要，因为它写在文档中:

Gradient scaling helps prevent gradients with small magnitudes from flushing to zero (“underflowing”) when training with mixed precision.

scaler = torch.cuda.amp.GradScaler()
for epoch in range(1):
    for input, target in zip(data, targets):
        with torch.cuda.amp.autocast():
            output = net(input)
            loss = loss_fn(output, target)

        scaler.scale(loss).backward()
        scaler.step(opt)
        scaler.update()
        opt.zero_grad()

还有，看NVIDIA Apex Documentation for PyTorch ，他们将其用作，

from apex import amp

model, optimizer = amp.initialize(model, optimizer)

loss = criterion(…)
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
optimizer.step()

我认为这也是 GradScaler 所做的，所以我认为这是必须的。有人可以在这里帮我查询吗。

Answer 1

简短回答:是的，如果没有 GradScaler()，您的模型可能无法收敛。

使用 FP16 存在三个基本问题:

• 权重更新:半精度，1 + 0.0001 舍入为 1。autocast() 负责这一项。
• 梯度消失:与单精度 2e-126 相比，使用半精度，任何小于(大致)2e-14 的值都舍入为 0。 GradScaler() 负责这个。
• 爆炸损失:与上述类似，半精度溢出的可能性也更大。这也由 autocast() 上下文管理。