4. 动态损失缩放策略调优

任务目标

在混合精度训练中，损失缩放（Loss Scaling）用于防止梯度下溢。本案例通过修改动态损失缩放的三个关键超参数——初始缩放值、增长因子与衰减因子，观察其对训练稳定性（是否出现NaN/Inf）与收敛速度（损失下降曲线）的影响，并确定一组较优参数。

环境与依赖

需要：Python 3.8+，PyTorch 1.12+（含torch.cuda.amp），CUDA可用GPU。若无可用的GPU环境，可改用CPU模式运行（但无法真实体现混合精度收益，仅用于验证逻辑正确性）。

核心实现

以下代码构建一个简单的二分类任务，使用torch.cuda.amp.GradScaler进行动态损失缩放训练。通过修改init_scale、growth_factor、backoff_factor三个参数，对比不同配置下的训练表现。

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

# 固定随机种子保证可复现
torch.manual_seed(42)

# 生成简单二分类数据
N = 2048
X = torch.randn(N, 16)
y = (X[:, 0] * 2 + X[:, 1] * 1.5 + torch.randn(N) * 0.1 > 0).long()

# 定义模型
model = nn.Sequential(
    nn.Linear(16, 64),
    nn.ReLU(),
    nn.Linear(64, 2)
)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 动态损失缩放配置（可调参数）
configs = [
    {"init_scale": 2**16, "growth_factor": 2.0, "backoff_factor": 0.5},   # 默认
    {"init_scale": 2**8,  "growth_factor": 2.0, "backoff_factor": 0.5},   # 低初始值
    {"init_scale": 2**24, "growth_factor": 2.0, "backoff_factor": 0.5},   # 高初始值
    {"init_scale": 2**16, "growth_factor": 1.5, "backoff_factor": 0.5},   # 慢增长
    {"init_scale": 2**16, "growth_factor": 4.0, "backoff_factor": 0.25},  # 快增长/深衰减
]

def train_with_config(cfg, epochs=30):
    # 重置模型与优化器
    model.apply(lambda m: m.reset_parameters() if hasattr(m, 'reset_parameters') else None)
    optimizer = optim.SGD(model.parameters(), lr=0.1)
    scaler = GradScaler(init_scale=cfg["init_scale"],
                        growth_factor=cfg["growth_factor"],
                        backoff_factor=cfg["backoff_factor"])
    losses = []
    nan_occurred = False
    for epoch in range(epochs):
        model.train()
        optimizer.zero_grad()
        with autocast():
            outputs = model(X)
            loss = criterion(outputs, y)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()
        losses.append(loss.item())
        if torch.isnan(loss) or torch.isinf(loss):
            nan_occurred = True
            break
    return losses, nan_occurred

# 运行所有配置并输出结果
if __name__ == "__main__":
    for i, cfg in enumerate(configs):
        losses, nan = train_with_config(cfg)
        final_loss = losses[-1] if losses else float('inf')
        status = "NaN/Inf" if nan else "稳定"
        print(f"配置{i+1}: init={cfg['init_scale']}, growth={cfg['growth_factor']}, "
              f"backoff={cfg['backoff_factor']} → {status}, 最终loss={final_loss:.4f}, "
              f"完成epoch数={len(losses)}")

运行与验证

运行方式：将上述代码保存为loss_scaling_tuning.py，在具备GPU的环境中执行python loss_scaling_tuning.py。若在CPU环境运行，需将autocast()替换为普通前向计算（但此时GradScaler不会生效，仅用于验证代码逻辑）。

验证判定标准：

• 观察输出中每种配置的最终loss与完成epoch数：若某配置出现NaN/Inf，说明该参数组合导致训练不稳定（梯度溢出或缩放不当）。
• 对比各配置的最终loss：loss越低且未中断，说明收敛越好。默认配置（配置1）通常能稳定收敛；若配置2（低初始值）在训练初期出现loss增大或NaN，说明初始缩放值过小导致梯度下溢；若配置3（高初始值）出现NaN，说明初始值过大导致梯度上溢。
• 配置4（慢增长）与配置5（快增长/深衰减）用于观察缩放因子调整速度的影响：增长过慢可能使缩放值长期偏小，衰减过深可能在遇到溢出时过度降低缩放值，均可能减缓收敛。

预期结果说明：由于本任务数据与模型固定，且使用固定随机种子，各配置的loss曲线是确定性的。配置1应稳定收敛至loss约0.1～0.3；配置2若出现NaN则说明初始缩放不足；配置3若出现NaN则说明初始缩放过大。实际输出以运行环境为准，上述为基于代码逻辑的预期推断，非实测数据。

调优要点

• 初始缩放值：应足够大以避免梯度下溢，但过大会导致上溢。常用范围2^16～2^24。
• 增长因子：默认2.0，表示连续无溢出时每步将缩放值乘以该因子。增大可更快达到合适范围，但可能跳过最优值。
• 衰减因子：默认0.5，表示发生溢出时将缩放值乘以该因子。减小可更激进地降低缩放值以恢复稳定，但会损失精度。
• 实际调参时应结合loss曲线与梯度统计：若频繁出现NaN，应降低初始值或增大衰减因子；若loss下降缓慢且无溢出，可尝试增大初始值或增长因子。