特殊混合精度示例

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

class MixedPrecisionModel(nn.Module):
    def __init__(self):
        super(MixedPrecisionModel, self).__init__()
        # 定义各层
        self.fp16_layer = nn.Linear(256, 128).half()  # 将这一层设为fp16精度
        self.mixed_layer = nn.Linear(128, 64)         # 这一层使用混合精度
        self.final_layer = nn.Linear(64, 10)          # 另一层使用混合精度

    def forward(self, x):
        # 将输入转换为fp16精度
        x = x.half()

        # fp16计算
        x = self.fp16_layer(x)

        # 使用autocast进行混合精度计算
        with autocast():
            x = self.mixed_layer(x)
            x = self.final_layer(x)
        
        return x

# 创建模型和数据
model = MixedPrecisionModel().cuda()
optimizer = torch.optim.Adam([
    {"params": model.fp16_layer.parameters(), "lr": 1e-3, "is_fp16": True},  # 对fp16层进行特殊设置
    {"params": model.mixed_layer.parameters()},
    {"params": model.final_layer.parameters()}
], lr=1e-3)
scaler = GradScaler()  # 创建梯度缩放器

# 创建一个示例输入数据
input_data = torch.randn(32, 256).cuda()  # 批大小32，输入维度256
target = torch.randint(0, 10, (32,)).cuda()


optimizer_fp16 = torch.optim.Adam(model.fp16_layer.parameters(), lr=1e-3)
optimizer_mixed = torch.optim.Adam(list(model.mixed_layer.parameters()) + list(model.final_layer.parameters()), lr=1e-3)

for epoch in range(10):
    optimizer_fp16.zero_grad()
    optimizer_mixed.zero_grad()
    
    # 启用混合精度训练
    with autocast():
        output = model(input_data)
        loss = nn.CrossEntropyLoss()(output, target)
    
    # 缩放损失反向传播并优化
    scaler.scale(loss).backward()
    
    # 更新混合精度层参数
    scaler.step(optimizer_mixed)
    scaler.update()
    
    # 手动更新FP16层参数
    for param in model.fp16_layer.parameters():
        if param.grad is not None:
            param.data -= optimizer_fp16.param_groups[0]["lr"] * param.grad
    
    print(f"Epoch {epoch+1}, Loss: {loss.item()}")