pytorch混合精度训练

# coding=utf-8
import torch
from torch.cuda.amp import autocast,GradScaler

def train(epoch):
    
    #梯度缩放 https://pytorch.org/docs/stable/amp.html#gradient-scaling
    scaler = GradScaler()

    x = torch.randn(3,100).cuda()
    y = torch.randn(3,5).cuda()
    model = torch.nn.Linear(100, 5).cuda()
    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
    
    for e in range(epoch):
        
        #充当上下文管理器或修饰器,混合精度运行 https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.autocast
        with autocast():
            y_pred = model(x)
            loss = torch.nn.functional.mse_loss(y_pred, y)

        #scale():将张量或张量列表乘以（“缩放”）比例因子。返回缩放的输出。
        #如果GradScaler未实例化，返回的输出将保持不变。
        scaler.scale(loss).backward()

        #step():内部调用unscale_(optimizer)将优化器的梯度张量除以比例因子。
        scaler.step(optimizer)

        # 更新比例因子
        scaler.update()

    print("complish")

if __name__=="__main__":
    if torch.cuda.is_available():
        train()
    else:
        print("cuda is not available")