torch.amp自动混合精度包

一、定义

1. 定义
2. 案例

二、实现

1. 定义
   torch.amp 给用户提供了较为方便的混合精度训练机制。方便”体现在两个方面：

   1. 用户不需要手动对模型参数 dtype 转换，amp 会自动为算子选择合适的数值精度
   2. 对于反向传播的时候，FP16 的梯度数值溢出的问题，amp 提供了梯度 scaling 操作，而且在优化器更新参数前，会自动对梯度 unscaling，所以，对用于模型优化的超参数不会有任何影响。
   3. 自动使用tensor Score 机制，加速训练。
   4. 组件：amp.autocast、amp.GradScaler.

2. 案例

from torch import autocast as autocast
model=Net().cuda()
optimizer=optim.SGD(model.parameters(),...)
scaler = GradScaler() #训练前实例化一个GradScaler对象
for epoch in epochs:
  for input,target in data:
    optimizer.zero_grad()

    with autocast(device_type="cuda"): ＃前后开启autocast
      output=model(input)
      loss = loss_fn(output,targt)

    scaler.scale(loss).backward()  #为了梯度放大
    #scaler.step() 首先把梯度值unscale回来，如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重，否则，忽略step调用，从而保证权重不更新。   
    scaler.step(optimizer)
    scaler.update()  #准备着，看是否要增大scaler

混合精度+梯度裁剪

from torch  import autocast as autocast

model=Net().cuda()
optimizer=optim.SGD(model.parameters(),...)

scaler = GradScaler() #训练前实例化一个GradScaler对象
for epoch in epochs:
  for input,target in data:
    optimizer.zero_grad()

    with autocast(): ＃前后开启autocast
      output=model(input)
      loss = loss_fn(output,targt)
    scaler.scale(loss).backward()  #为了梯度放大
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)  
    scaler.step(optimizer)
    scaler.update()  #准备着，看是否要增大scaler

数据并行+混合精度

MyModel(nn.Module):
    def forward(self, input):
        with autocast():
            ...
model = MyModel()
dp_model=nn.DataParallel(model)
with autocast():
    output=dp_model(input)
    loss = loss_fn(output)