混合精度训练（Mixed Precision Training）

# FP32主权重
master_weights = model.parameters().float()

# FP16计算权重
compute_weights = model.half()

# 更新时使用FP32
optimizer.step()  # 更新master_weights
compute_weights.copy_(master_weights.half())

# 初始化高损失缩放值（如2^24）
loss_scale = 2**24

# 训练循环
loss = model(input) * loss_scale
loss.backward()

if has_overflow(gradients):
    # 梯度溢出，减半缩放值
    loss_scale /= 2
    skip_optimizer_step()
else:
    # 正常更新，尝试增加缩放值
    unscale_gradients(loss_scale)
    optimizer.step()
    if steps % 2000 == 0:
        loss_scale *= 2

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for data, target in dataloader:
    optimizer.zero_grad()
    
    # 自动混合精度区域
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    
    # 损失缩放和反向传播
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()