GradScaler MaxClipGradScaler

本文是不求甚解，抛转引玉，欢迎大家发表更好的意见，

insightface训练代码：

个人感觉先MaxClipGradScaler，再：

# 梯度裁剪 , 求所有参数的二范数，如果大于max_norm ，都乘以　max_norm/所有参数的二范数
 clip_grad_norm_(backbone.parameters(), max_norm=5, norm_type=2)

grad_scaler = MaxClipGradScaler(cfg.batch_size, 128 * cfg.batch_size, growth_interval=100) if cfg.fp16 else None
    for epoch in range(start_epoch, cfg.num_epoch):
        train_sampler.set_epoch(epoch)
        for step, (img, label) in enumerate(train_loader):
            global_step += 1
            features = F.normalize(backbone(img))
            x_grad, loss_v = module_partial_fc.forward_backward(label, features, opt_pfc, backbone)
            if cfg.fp16:
                features.backward(grad_scaler.scale(x_grad))
                grad_scaler.unscale_(opt_backbone)
                clip_grad_norm_(backbone.parameters(), max_norm=5, norm_type=2)
                grad_scaler.step(opt_backbone)
                grad_scaler.update()
            else:
                features.backward(x_grad)
                # 梯度裁剪 , 求所有参数的二范数，如果大于max_norm ，都乘以　max_norm/所有参数的二范数
                clip_grad_norm_(backbone.parameters(), max_norm=5, norm_type=2)
                opt_backbone.step()
 
            opt_pfc.step()
            module_partial_fc.update()
            opt_backbone.zero_grad()
            opt_pfc.zero_grad()
            loss.update(loss_v, 1)
            callback_logging(global_step, loss, epoch, cfg.fp16, grad_scaler)
            callback_verification(global_step, backbone)
        callback_checkpoint(global_step, backbone, module_partial_fc)
        scheduler_backbone.step()
        scheduler_pfc.step()
    dist.destroy_process_group()

以下内容转自：

Pytorch自动混合精度(AMP)介绍与使用 - jimchen1218 - 博客园

2)损失放大（Loss scaling)

      即使了混合精度训练，还是存在无法收敛的情况，原因是激活梯度的值太小，造成了溢出。可以通过使用torch.cuda.amp.GradScaler，通过放大loss的值来防止梯度的underflow（只在BP时传递梯度信息使用，真正更新权重时还是要把放大的梯度再unscale回去）；

反向传播前，将损失变化手动增大2^k倍，因此反向传播时得到的中间变量（激活函数梯度）则不会溢出；

反向传播后，将权重梯度缩小2^k倍，恢复正常值。

三．如何使用AMP?

　目前有两种版本：pytorch1.5之前使用的NVIDIA的三方包apex.amp和pytorch1.6自带的torch.cuda.amp

１．pytorch1.5之前的版本（包括１．５）

 使用方法如下：

from apex import amp
model,optimizer = amp.initial(model,optimizer,opt_level="O1")   #注意是O,不是０
with amp.scale_loss(loss,optimizer) as scaled_loss:
    scaled_loss.backward()
取代
loss.backward()

其中,opt_level配置如下：

　O0:纯FP32训练，可作为accuracy的baseline；

　O1:混合精度训练（推荐使用），根据黑白名单自动决定使用FP16(GEMM,卷积）还是FP32（softmax)进行计算。

　O2:几乎FP16，混合精度训练，不存在黑白名单　，除了bacthnorm，几乎都是用FP16计算；

　O3:纯FP16训练，很不稳定，但是可以作为speed的baseline；

动态损失放大（dynamic loss scaling)部分，为了充分利用FP16的范围，缓解舍入误差，尽量使用最高的放大倍数2^24,如果产生上溢出，则跳出参数更新，缩小放大倍数使其不溢出。在一定步数后再尝试使用大的scale来充分利用FP16的范围。

 

２）GradScaler

　　使用前，需要在训练最开始前实例化一个GradScaler对象，例程如下：

from torch.cuda.amp 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.step()　首先把梯度值unscale回来，如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重，否则，忽略step调用，从而保证权重不更新。
　　 scaler.step(optimizer)
    scaler.update()  #准备着，看是否要增大scaler

　　scaler的大小在每次迭代中动态估计，为了尽可能减少梯度underflow，scaler应该更大；但太大，半精度浮点型又容易overflow（变成inf或NaN).所以，动态估计原理就是在不出现if或NaN梯度的情况下，尽可能的增大scaler值。在每次scaler.step(optimizer)中，都会检查是否有inf或NaN的梯度出现：

　　１．如果出现inf或NaN,scaler.step(optimizer)会忽略此次权重更新(optimizer.step()），并将scaler的大小缩小（乘上backoff_factor)；

　　２．如果没有出现inf或NaN,那么权重正常更新，并且当连续多次(growth_interval指定)没有出现inf或NaN，则scaler.update()会将scaler的大小增加(乘上growth_factor)。