GradScaler MaxClipGradScaler

x33g5p2x  于2022-06-24 转载在 其他  
字(3.0k)|赞(0)|评价(0)|浏览(493)

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

insightface训练代码:

个人感觉先MaxClipGradScaler,再:

梯度裁剪 , 求所有参数的二范数,如果大于max_norm ,都乘以 max_norm/所有参数的二范数

 clip_grad_norm_(backbone.parameters(), max_norm=5, norm_type=2)

  1. grad_scaler = MaxClipGradScaler(cfg.batch_size, 128 * cfg.batch_size, growth_interval=100) if cfg.fp16 else None
  2. for epoch in range(start_epoch, cfg.num_epoch):
  3. train_sampler.set_epoch(epoch)
  4. for step, (img, label) in enumerate(train_loader):
  5. global_step += 1
  6. features = F.normalize(backbone(img))
  7. x_grad, loss_v = module_partial_fc.forward_backward(label, features, opt_pfc, backbone)
  8. if cfg.fp16:
  9. features.backward(grad_scaler.scale(x_grad))
  10. grad_scaler.unscale_(opt_backbone)
  11. clip_grad_norm_(backbone.parameters(), max_norm=5, norm_type=2)
  12. grad_scaler.step(opt_backbone)
  13. grad_scaler.update()
  14. else:
  15. features.backward(x_grad)
  16. # 梯度裁剪 , 求所有参数的二范数,如果大于max_norm ,都乘以 max_norm/所有参数的二范数
  17. clip_grad_norm_(backbone.parameters(), max_norm=5, norm_type=2)
  18. opt_backbone.step()
  19. opt_pfc.step()
  20. module_partial_fc.update()
  21. opt_backbone.zero_grad()
  22. opt_pfc.zero_grad()
  23. loss.update(loss_v, 1)
  24. callback_logging(global_step, loss, epoch, cfg.fp16, grad_scaler)
  25. callback_verification(global_step, backbone)
  26. callback_checkpoint(global_step, backbone, module_partial_fc)
  27. scheduler_backbone.step()
  28. scheduler_pfc.step()
  29. 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

1.pytorch1.5之前的版本(包括1.5)

使用方法如下:

  1. from apex import amp
  2. model,optimizer = amp.initial(model,optimizer,opt_level="O1") #注意是O,不是0
  3. with amp.scale_loss(loss,optimizer) as scaled_loss:
  4. scaled_loss.backward()
  5. 取代
  6. 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的范围。

2)GradScaler

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

  1. from torch.cuda.amp import autocast as autocast
  2. model=Net().cuda()
  3. optimizer=optim.SGD(model.parameters(),...)
  4. scaler = GradScaler() #训练前实例化一个GradScaler对象
  5. for epoch in epochs:
  6. for input,target in data:
  7. optimizer.zero_grad()
  8. with autocast(): #前后开启autocast
  9. output=model(input)
  10. loss = loss_fn(output,targt)
  11. scaler.scale(loss).backward() #为了梯度放大
  12. #scaler.step() 首先把梯度值unscale回来,如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重,否则,忽略step调用,从而保证权重不更新。
  13.    scaler.step(optimizer)
  14. scaler.update() #准备着,看是否要增大scaler

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

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

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

相关文章