GradScaler是否需要与pytorch进行混合精确训练?

Question

所以AMP:自动混合精度训练教程Normal网络上，我发现有两个版本，Automatic和GradScaler。我只是想知道这是否明智/必要在培训中使用GradScaler，因为文档中写着:

梯度缩放有助于防止在混合精度训练时小幅度的梯度冲洗为零("下流")。

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的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-14轮到0，而不是单精度2e-126。GradScaler()负责这个
• 爆炸损失:与上面类似，半精度时溢出的可能性也大得多。这也由autocast()上下文管理。