cuda.amp 和 model.half() 有什么区别？

Question

根据https://pytorch.org/blog/acceleating-training-on-nvidia-gpus-with-pytorch-automatic-mixed- precision/

我们可以用：

   with torch.cuda.amp.autocast():
      loss = model(data)

为了将操作转换为混合精度。

另一件事是我们可以用来 model.half()将所有模型权重转换为半精度。

1. 这两个命令有什么区别？
2. 如果我想利用FP16（为了创建更大的模型和更短的训练时间），我需要什么？我需要使用 model.half()还是正在使用torch.cuda.amp（根据上面的链接）？

Answer 1

如果将整个模型转换为 fp16，某些激活函数和批归一化层可能会导致 fp16 权重下溢，即变为零。因此，始终建议使用 autocast，它会在有问题的层中将权重内部转换为 fp32。

model.half()最终将在 FP16 中节省重量，而自动施放权重仍将在 FP32 中。在 FP16 中训练会比自动施法更快，但如果不小心的话，不稳定的可能性会更高。使用自动投射时，您还需要在反向传播期间放大梯度

如果推理端需要 fp16，我建议使用 autocast，然后使用 ONNX 和 tensorrt 转换为 fp16。