当前位置：首页 > news >正文

PyTorch-混合精度训练（amp）

news 2025/11/15 12:04:02

1. 核心目标：为什么使用AMP？

在我第一次使用AMP（auto Mixed Precision，amp）时，我观察到训练速度显著提升，同时解决了在mit-states数据集上的显存不足问题。AMP的核心目标就是实现这两点：

加速训练：在现代NVIDIA GPU（如A100）上，利用其内置的Tensor Cores硬件单元，对float16（半精度）的计算速度远超float32（单精度）。
减少显存占用：float16数据类型占用的显存空间是float32的一半。这意味着模型参数、梯度和中间激活值所占用的显存都能大幅降低，从而允许我们使用更大的模型或更大的批次大小（batch size）。

2. AMP的两大核心组件与工作原理

PyTorch通过torch.cuda.amp模块提供了两个核心工具来实现自动混合精度，分别是autocast和GradScaler。

a. autocast：智能的自动类型转换

autocast是一个上下文管理器，它会自动为不同类型的运算选择合适的精度，实现“混合精度”。

工作原理：

对性能影响大的运算 → float16：对于大部分计算密集型操作，如矩阵乘法（nn.Linear）和卷积（nn.Conv2d），autocast会将它们的输入和计算过程切换到float16，以充分利用Tensor Cores进行加速。
对精度要求高的运算 → float32：对于一些可能因精度损失而导致数值不稳定的操作，如损失函数计算、BatchNorm的统计量更新等，autocast会智能地将它们保持在float32下进行，以确保模型的准确性和稳定性。

在我的代码中 (train_autotest.py):

我将整个前向传播和损失计算过程都包裹在了with autocast():代码块中。

Python# train_autotest.py from torch.cuda.amp import autocast  
# ... 在训练循环中 ... 
with autocast():     predict = model(batch, train_pairs)     loss = model.loss_calu(predict, batch)     loss = loss / config.gradient_accumulation_steps

这确保了模型内部的运算能够被自动、智能地分配到最合适的精度上执行。

b. GradScaler：防止梯度下溢的“保险丝”

这是成功进行混合精度训练的关键保障。

问题背景：梯度下溢 (Gradient Underflow)
float16的数值表示范围比float32小得多。在反向传播过程中，计算出的梯度可能非常小，以至于超出了float16能表示的最小范围，从而被舍入为0。如果大量梯度变成0，模型的参数就无法得到有效更新，训练就会失败。
工作原理:
1.放大损失 (Scaling)：在调用.backward()之前，GradScaler会将计算出的损失值乘以一个巨大的缩放因子（例如2^16）。
2.放大梯度 (Chain Rule)：根据链式法则，损失被放大后，所有计算出的梯度也会被同等放大。这样，原本微小的梯度就被“抬升”到了float16可以安全表示的范围内，避免了下溢问题。
3.缩回梯度 (Unscaling)：在优化器更新权重之前（调用optimizer.step()时），GradScaler会先将梯度“缩放”回其原始大小，确保权重更新的步长是正确的。
4.动态调整: GradScaler还会动态监测梯度中是否出现inf或NaN（这可能在缩放后发生），如果出现，它会自动跳过当次的权重更新，并在下一次迭代中减小缩放因子，以保持训练的稳定性。
在我的代码中 (train_autotest.py):

我严格按照AMP的标准流程使用了GradScaler。

Python# train_autotest.py from torch.cuda.amp import GradScaler  
# 1. 在训练开始前初始化 
scaler = GradScaler()  # ... 在训练循环中 ... 
# 2. 用scaler放大损失并进行反向传播 
scaler.scale(loss).backward()  # ... 在累积了足够的梯度后 ... 
# 3. 用scaler.step()来unscale梯度并更新权重 
scaler.step(optimizer) # 4. 更新scaler的缩放因子 
scaler.update() 
optimizer.

3. 总结与最佳实践

适用场景: AMP最适用于现代NVIDIA GPU（Volta架构及以后，如V100, T4, A100, H100等），因为这些GPU配备了Tensor Cores。
对性能的影响: 对于绝大多数基于Transformer和卷积的现代模型，AMP对最终的模型精度影响极小，可以认为是“性能无损”的。它带来的训练加速和显存节省是巨大的工程收益，可以显著加快研究和迭代速度。
代码实现: 必须同时使用autocast和GradScaler。只使用autocast而不使用GradScaler会导致梯度下溢，很可能使训练失败或效果变差。
评估阶段: 在evaluate或test函数中，通常也建议使用with autocast():来包裹模型的前向传播，这样可以加速推理过程，且通常不会影响最终的评估结果。我的test_autotest.py代码中也遵循了这一点。

参考

PyTorch自动混合精度

PyTorch对混合精度的支持始于1.6版本，位于torch.cuda.amp模块下，主要是torch.cuda.amp.autocast和torch.cuda.amp.GradScale两个模块，autocast针对选定的代码块自动选取适合的计算精度，以便在保持模型准确率的情况下最大化改善训练效率；GradScaler通过梯度缩放，以最大程度避免使用FP16进行运算时的梯度下溢。官方给的使用这两个模块进行自动精度训练的示例代码链接给出，我对其示例解析如下，这就是一般的训练框架。

# 以默认精度创建模型和优化器
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)# 创建梯度缩放器
scaler = GradScaler()for epoch in epochs:for input, target in data:optimizer.zero_grad()# 通过自动类型转换进行前向传播with autocast():output = model(input)loss = loss_fn(output, target)# 缩放大损失，反向传播不建议放到autocast下，它默认和前向采用相同的计算精度scaler.scale(loss).backward()# 先反缩放梯度，若反缩后梯度不是inf或者nan，则用于权重更新scaler.step(optimizer)# 更新缩放器scaler.update()