【深度学习 混合精度训练】只看训练过程就够了

传统训练（FP32全精度）：
输入数据 → 模型计算 → 损失计算 → 反向传播 → 更新参数
   ↓           ↓          ↓           ↓          ↓
 FP32        FP32       FP32        FP32       FP32
（稳但慢） （稳但慢）  （稳但慢）   （稳但慢）  （稳但慢）

混合精度训练（FP16+FP32）：
输入数据 → 模型计算 → 损失计算 → 反向传播 → 更新参数
   ↓           ↓          ↓           ↓          ↓
 FP16        FP16       FP16        FP16       FP32
（快点）    （快2-3倍） （快点）    （快点）   （这里必须用FP32！）
         ↑                        ↑
     用Tensor Core加速       损失要放大再传回去

# 原来正常的训练循环
for data, target in dataloader:
    optimizer.zero_grad()
    output = model(data)          # 前向
    loss = criterion(output, target)
    loss.backward()               # 反向
    optimizer.step()              # 更新

# 启用混合精度后（只改4处！）
from torch.cuda.amp import autocast, GradScaler  # 1. 导入

scaler = GradScaler()  # 2. 创建缩放器

for data, target in dataloader:
    optimizer.zero_grad()
    
    with autocast():  # 3. 包住前向计算（自动用FP16算）
        output = model(data)
        loss = criterion(output, target)
    
    scaler.scale(loss).backward()  # 4. 用scaler包装loss
    scaler.step(optimizer)         #   用scaler包装step
    scaler.update()                #   更新缩放因子

# 假设梯度是 0.000001（FP16存不了，会变成0）
# GradScaler做的事：
1. 把损失放大：loss * 1024 = 变大的loss
2. 反向传播：得到放大的梯度
3. 更新前缩小：梯度 / 1024 = 正确的梯度
4. 用FP32更新参数

# 原来：batch_size=32, lr=0.001
# 混合精度后：batch_size=64, lr=0.002
# 并且加学习率热身

# 正常情况：
# 损失稳步下降 → 训练正常
# 准确率逐步上升 → 训练正常

# 有问题的情况：
# 损失突然变成NaN → 梯度爆炸了（调小学习率）
# 损失完全不变 → 梯度消失了（检查缩放器）

# 运行训练时，另一个终端运行：
watch -n 1 nvidia-smi

# 应该看到：
# - GPU内存占用减少40-50%
# - GPU利用率更高

# 每个epoch结束后：
model.eval()  # 重要！验证时不要用autocast
with torch.no_grad():  # 重要！不要计算梯度
    for data, target in val_loader:
        output = model(data)  # 这里用FP32算
        # 计算准确率...

# 1. 降低学习率（最有效）
optimizer = Adam(model.parameters(), lr=0.001)  # 原来是0.002

# 2. 改缩放器设置
scaler = GradScaler(init_scale=65536.0)  # 默认是65536.0

# 3. 加梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 检查是否是这些问题：
1. 数据加载是瓶颈（用DataLoader num_workers>0）
2. 模型有些层不支持FP16（如某些自定义层）
3. GPU太旧（需要Pascal架构以上）

# 1. 试试BF16（如果GPU支持）
# 训练命令加：--bf16

# 2. 调整优化器
optimizer = AdamW(model.parameters(), lr=0.002, weight_decay=0.05)

# 3. 延长训练时间（有时需要多练几个epoch）

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

def train_one_epoch(model, dataloader, optimizer, epoch):
    """一个epoch的训练"""
    model.train()
    scaler = GradScaler()  # 每个epoch都可以新建
    
    for batch_idx, (data, target) in enumerate(dataloader):
        data, target = data.cuda(), target.cuda()
        
        # 1. 清空梯度
        optimizer.zero_grad()
        
        # 2. 前向计算（自动混合精度）
        with autocast():
            output = model(data)
            loss = nn.CrossEntropyLoss()(output, target)
        
        # 3. 反向传播（自动梯度缩放）
        scaler.scale(loss).backward()
        
        # 4. 参数更新（自动取消缩放）
        scaler.step(optimizer)
        scaler.update()
        
        # 5. 打印进度
        if batch_idx % 100 == 0:
            print(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}')
    
    return loss.item()

# 使用：直接调用这个函数训练