PyTorch 中的混合精度训练方法，从 autocast 到 GradScalar

PyTorch 的混合精度训练主要由两个方法实现：amp.autocast 和 amp.GradScalar。在这两个工具的帮助下，可以实现以 torch.float16 的混合精度训练。当然，这两个方法都是模块化并且通常都会一起调用，但并不一定总是需要一起使用。

参考：

Automatic Mixed Precision package - torch.amp
Automatic Mixed Precision examples
Automatic Mixed Precision recipe

一 amp.autocast

在 PyTorch 中，autocast 是一个用于自动混合精度训练的上下文管理器。它的主要作用是根据操作的类型自动选择使用 float16（半精度）或 float32（单精度）进行计算，从而提高训练速度并减少内存使用。以下是 autocast 的具体功能和工作原理：

1. 自动选择精度
   操作类型：在深度学习中，不同的操作对数值精度的要求不同。例如，某些操作（如矩阵乘法）在 float16 下可以安全地执行，而其他操作（如某些归一化或损失计算）可能需要 float32 来保持数值稳定性。
   自动化：使用 autocast 时，PyTorch 会自动判断每个操作的最佳精度，并在适当的情况下使用 float16，在其他情况下使用 float32。这使得开发者不需要手动管理每个操作的精度，从而简化了代码。
2. 上下文管理器
   使用方式：autocast 通常用作上下文管理器，使用 with torch.cuda.amp.autocast(): 语句包裹需要进行混合精度计算的代码块。在这个代码块内，所有的张量操作都会根据 autocast 的规则自动选择精度。
3. 性能提升
   加速训练：通过使用 float16 进行计算，autocast 可以显著提高训练速度，尤其是在支持半精度计算的 GPU 上（如 NVIDIA 的 Volta 和 Ampere 架构）。减少内存使用：使用 float16 还可以减少显存的占用，使得更大的模型或更大的批量大小成为可能。
4. 示例
   一般来说，建议在模型前向传递和 loss计算中使用 autocast，而不建议在反向传播和参数更新过程中使用，例如：

# 来源：https://pytorch.org/docs/stable/amp.html#torch.autocast
# Creates model and optimizer in default precision
model = Net(