混合精度训练调优：不同学习率策略对AMP效果的影响

混合精度训练调优：不同学习率策略对AMP效果的影响

最近在做PyTorch模型优化时，遇到了一个很实际的问题：混合精度训练（AMP）的效果竟然和学习率策略强相关。这里记录一下踩坑过程。

问题背景

使用torch.cuda.amp.GradScaler进行AMP训练时，发现不同学习率策略下，模型收敛速度和最终精度差异巨大。原本以为AMP只是加速训练，没想到对学习率调优也有影响。

实验设置

import torch
import torch.nn as nn
from torch.cuda.amp import GradScaler, autocast
from torch.optim.lr_scheduler import StepLR, CosineAnnealingLR

# 模型结构
model = nn.Sequential(
    nn.Linear(784, 512),
    nn.ReLU(),
    nn.Linear(512, 10)
).cuda()

# 优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
scaler = GradScaler()

三种学习率策略对比

策略1：固定学习率（lr=0.01）

scheduler = None
# 训练循环中使用 amp
for epoch in range(5):
    for batch in dataloader:
        optimizer.zero_grad()
        with autocast():
            output = model(batch)
            loss = criterion(output, target)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

策略2：StepLR衰减（gamma=0.1）

scheduler = StepLR(optimizer, step_size=2, gamma=0.1)
# 训练循环保持不变

策略3：CosineAnnealingLR

scheduler = CosineAnnealingLR(optimizer, T_max=5)
# 训练循环保持不变

实验结果（训练5个epoch）

踩坑总结

1. AMP对学习率敏感：固定LR时AMP效果最差，CosineAnnealing效果最好
2. 学习率衰减策略影响梯度更新幅度，进而影响AMP的loss scaling
3. 建议：使用AMP时配合动态学习率策略，并在训练初期使用较大学习率

实践建议

• 优先尝试CosineAnnealingLR配合AMP
• 调整初始学习率为0.05~0.1之间
• 使用torch.cuda.amp.GradScaler().scale()来检查缩放因子变化