PyTorch ： torch.cuda.amp: 自动混合精度详解

amp : 全称为 Automatic mixed precision，自动混合精度

amp总体介绍：

可以在神经网络推理过程中，针对不同的层，采用不同的数据精度进行计算，从而实现节省显存和加快速度的目的。

通常，深度学习中使用的精度为32位（单精度）浮点数，而使用16位（半精度）浮点数可以将内存使用减半，同时还可以加快计算速度。然而，16位浮点数的精度较低，可能导致数值下溢或溢出，从而影响训练结果。

混合精度: 有不止一种精度的Tensor（不同精度的数值计算混合使用来加速训练和减少显存占用） :

• torch.FloatTensor（浮点型 32位）(torch默认的tensor精度类型是torch.FloatTensor)
• torch.HalfTensor（半精度浮点型 16位）

自动：

• 预示着Tensor的dtype类型会自动变化，也就是框架按需自动调整tensor的dtype

使用自动混合精度 (amp) 的原因

• torch.HalfTensor的优势是存储小、计算快、更好的利用CUDA设备的Tensor Core。因此训练的时候可以减少显存的占用（可以增加batchsize了），同时训练速度更快
• torch.HalfTensor的劣势是数值范围小（更容易Overflow /Underflow）、舍入误差（Rounding Error，导致一些微小的梯度信息达不到16bit精度的最低分辨率，从而丢失）

混合精度训练机制（API：autocast类；GradScaler 类`）

1. amp.autocast：

• amp.autocast()：用户不需要手动对模型参数 dtype 转换，amp 会自动为算子选择合适的数值精度；是PyTorch中一种混合精度的技术，可在保持数值精度的情况下提高训练速度和减少显存占用。
• amp.autocast()：能够自动将16位浮点数转换为32位浮点数进行数值计算，并在必要时将结果转换回16位浮点数。这种自动转换可以帮助避免数值下溢或溢出的问题，并在保持数值精度的同时提高计算速度和减少显存占用。
• autocast的优缺点
  使用autocast的优点在于，它可以自动选择合适的精度进行计算，从而提高了计算速度和内存使用效率，同时也减少了代码实现的复杂性。
  缺点在于，autocast并不是所有运算都能自动找到合适的精度进行计算，如果没有找到合适的精度，就会使用默认精度（FP32），这样就会增加显存的使用，并且在需要处理极大或极小的数据时，可能会出现数值精度问题。
• 使用torch.cuda.amp.autocast()的过程如下：

• 将模型和数据移动到GPU上
• 使用torch.cuda.amp.autocast()上下文管理器包装模型的前向传递和损失计算
• 使用scaler（即torch.cuda.amp.GradScaler对象）将反向传播的梯度缩放回16位
• 执行梯度更新

问：使用 torch.cuda.amp.autocast() 将数据 从32位（单精度） 转换为 16位（半精度），会导致精度丢失嘛？
答：使用 torch.cuda.amp.autocast() 将数据从32位（单精度）转换为16位（半精度）会导致精度损失。由于16位浮点数只能表示更少的有效位数，因此它们的精度不如32位浮点数。在混合精度训练中，为了平衡精度和性能，通常会将网络的前向传播和反向传播过程中的参数和梯度计算使用半精度浮点数来加速计算。这种方法可以在一定程度上降低计算精度要求，但会带来一定的精度损失。
·
尽管存在精度损失，使用半精度浮点数的优点在于它们可以显著降低计算时间和显存消耗，从而使模型可以在更大的批量下进行训练，提高训练效率。此外，在实际应用中，对于某些任务，半精度精度的计算误差对于结果的影响可能不是很大，因此，半精度计算可以在保证结果准确性的前提下，大幅度提高模型的训练速度和效率。
原文链接：https://blog.csdn.net/weixin_37804469/article/details/129733868

2. amp.GradScaler：

• 对于反向传播的时候，FP16 的梯度数值溢出的问题，amp 提供了梯度 scaling 操作，而且在优化器更新参数前，会自动对梯度unscaling，所以，对用于模型优化的超参数不会有任何影响；
• 具体来说，GradScaler 可以将梯度缩放到较小的范围，以避免数值下溢或溢出的问题，同时保持足够的精度以避免模型的性能下降。

2.1 GradScaler 类

torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)

其中：

• init_scale: scale factor 的初始值
• growth_factor: 每次 scale factor 的增长系数
• backoff_factor: scale factor 下降系数
• growth_interval: 每隔多个interval 增长 scale factor
• enabled: 是否做 scale

2.2 scale(output)方法

对outputs乘 scale factor，并返回，如果enabled=False就原样返回。

2.3 step(optimizer, *args, **kwargs)方法

step 方法在做两件事情：

1. 对梯度 unscale，如果之前没有手动调用unscale方法的话

2. 检查梯度溢出，如果没有nan/inf，就执行 optimizer 的 step，如果有就跳过

   注意：GradScaler的step不支持传 closure。

2.4 update(new_scale=None)方法

update方法在每个 iteration 结束前都需要调用，如果参数更新跳过，会给 scale factor 乘backoff_factor，或者到了该增长的 iteration，就给 scale factor 乘growth_factor。也可以用new_scale直接更新 scale factor。

以下是一个示例，展示了如何在 PyTorch 中使用 GradScaler：

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

# 创建 GradScaler 和模型
scaler = GradScaler()
model = torch.nn.Linear(10, 1).cuda()

# 定义损失函数和优化器
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

# 定义一些训练数据和目标
x = torch.randn(32, 10).cuda()
y = torch.randn(32, 1).cuda()

# 使用 GradScaler 进行自动混合精度训练
for i in range(1000):
    optimizer.zero_grad()

    # 将前向传递包装在autocast中以启用混合精度
    with autocast():
        y_pred = model(x)
        loss = loss_fn(y_pred, y)

    # 调用 GradScaler 的 backward() 方法计算梯度并缩放
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

    if i % 100 == 0:
        print(f"Step {i}, loss={loss.item():.4f}")

在这个示例中:

1. 创建一个 GradScaler 对象 scaler
2. 定义模型和优化器
3. 在训练循环中，使用 autocast() 上下文管理器将前向传递操作包装起来，这样就可以使用混合精度进行计算
4. 调用 scaler.scale(loss) 计算损失的缩放版本，并调用 scaler.step(optimizer) 来更新模型参数
5. 最后使用 scaler.update() 更新 GradScaler 对象的内部状态
6. 这个过程可以重复进行多次，直到训练结束。

原文链接：https://blog.csdn.net/qq_43369406/article/details/130393078