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().cuda()
optimizer = optim.SGD(model.parameters(), ...)

for input, target in data:
    optimizer.zero_grad()

    # Enables autocasting for the forward pass (model + loss)
    with torch.autocast(device_type="cuda"):
        output = model(input)　# 前向传递
        loss = loss_fn(output, target) # loss计算

    # Exits the context manager before backward()
    loss.backward() # 后向传递（计算梯度）
    optimizer.step() # 参数更新

5. 总结
   autocast 是 PyTorch 中实现混合精度训练的关键工具，它通过自动选择操作的精度来提高训练速度和减少内存使用，同时保持数值稳定性。使用 autocast 可以让开发者更专注于模型的设计和训练，而不必担心每个操作的精度管理。

二 amp.GradScalar

GradScaler 是 PyTorch 中用于混合精度训练的一个重要组件，主要用于处理梯度的缩放，以确保在使用半精度（float16）进行训练时的数值稳定性。以下是对 GradScaler 的详细介绍：

1. 背景
   在混合精度训练中，使用 float16 进行计算可以提高速度和减少内存使用，但由于 float16 的数值范围和精度较低，可能会导致梯度下溢（即梯度变得过小而被视为零）或上溢（即梯度变得过大而无法表示）。GradScaler 的作用就是解决这个问题。
2. 主要功能
   动态缩放：GradScaler 会在反向传播时动态调整损失值的缩放因子，以避免梯度的数值不稳定。它会根据当前的训练状态自动选择合适的缩放因子。
   防止下溢和上溢：通过将损失值乘以一个缩放因子，GradScaler 可以确保计算出的梯度不会因为数值过小而消失，也不会因为数值过大而溢出。
3. 使用流程
   使用 GradScaler 的典型流程如下：

初始化：创建 GradScaler 的实例;

scaler = torch.cuda.amp.GradScaler()

前向传播：在前向传播时，使用 autocast 来自动选择精度，不建议使用GradScaler；

   with torch.cuda.amp.autocast():
       output = model(inputs)
       loss = compute_loss(output, targets)

反向传播:

• 使用 scaler.scale(loss).backward() 进行反向传播。这里，scale 方法会将损失值缩放，并计算梯度。
• 使用 scaler.step(optimizer) 更新模型参数。此步骤会将缩放后的梯度应用到优化器中。
• 使用 scaler.update() 更新缩放因子，以便在下一次迭代中使用。

4. 示例代码

以下是一个使用 GradScaler 的简单示例：

import torch
from torch.cuda.amp import GradScaler, autocast

model = ...  # 初始化模型
optimizer = ...  # 初始化优化器
scaler = GradScaler()  # 创建 GradScaler 实例

for data, target in dataloader:
    optimizer.zero_grad()  # 清零梯度

    with autocast():  # 自动选择精度
        output = model(data) # 前向传递
        loss = compute_loss(output, target) # 计算loss

    scaler.scale(loss).backward()  # 缩放损失并反向传播以计算梯度
    scaler.step(optimizer)  # 更新模型参数
    scaler.update()  # 更新缩放因子

5. 总结
   GradScaler 是 PyTorch 中实现混合精度训练的关键工具，它通过动态缩放损失值来确保在使用 float16 进行训练时的数值稳定性。使用 GradScaler 可以有效地避免梯度下溢和上溢的问题，从而提高训练的可靠性和效率。通过结合 autocast 和 GradScaler，开发者可以在不牺牲模型性能的情况下，充分利用混合精度训练的优势。