pytorch + apex 生活变得更美好

2024年2月27日 436次阅读 来源: lbin

pytorch1.0发布以后，分布式训练变得异常顺滑，虽然apex很早就放出来了，似乎也没有了使用的必要。但其实apex作为一个补丁，其实也解决了几个重要的问题，加上以后，模型的训练会更得心应手。

pytorch当前存在的几个问题

第一个是fp16的问题，pytorch原生是可以把模型转换为fp16，但是训练的时候会产生很多模范，尤其是模型含有Batchnorm的时候。

model.half()
for layer in model.modules():
  if isinstance(layer, nn.BatchNorm2d):
    layer.float()

训练的时候需要将batchnorm层转为fp32, 然后这又涉及了混合精度训练的问题，无论怎样都会出现难以收敛的问题。

第二个问题是sync bn的问题，大家都知道在检测或者分割模型的训练中，sync bn能够极大加快模型的收敛，精度也会有所提升。然而原生的sync bn一直未能release，其他第三方的实现多多少少有些问题。

Apex

apex是由nvidia维护的一个支持混合精度分布式训练的第三方pytorch扩展。实测可以完美地解决上述两个问题

GitHub地址：https://github.com/NVIDIA/apex

2月底会release正式的1.0, 不过还有几个小时，所以不一定是放鸽子了。

混合精度训练

apex.FP16_Optimizer目前能够完整的复现fp32的训练结果，并且速度要更快，官方文档里用了“”highway” for FP16 training”

分布式训练

其中的亮点就是内置了SyncBatchNorm。

apex.parallel.SyncBatchNorm extends torch.nn.modules.batchnorm._BatchNorm to support synchronized BN. It reduces stats across processes during multiprocess distributed data parallel training. Synchronous Batch Normalization has been used in cases where only very small number of mini-batch could be fit on each GPU. All-reduced stats boost the effective batch size for sync BN layer to be the total number of mini-batches across all processes. It has improved the converged accuracy in some of our research models.

集成进pytorch

step 1

amp_handle = amp.init(enabled=True)

step 2

loss = criterion(…)
loss.backward()
optimizer.step()

step 3

loss = criterion(…)
with amp_handle.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
optimizer.step()

for FP16_Optimizer

optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)

过些时间放出mask-rcnn上添加apex的比较，敬请期待。

    原文作者：lbin
    原文地址: https://zhuanlan.zhihu.com/p/57958993
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。
点赞