Profiling数据发现Loss Scale在梯度更新引入大量Add、Mul小算子。
融合小算子、并替换亲和API减少梯度更新阶段的add、mul算子调用数量。
该环境变量在使用混合精度时,已默认提醒用户开启。
替换亲和函数。原始代码如下:
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum) model, optimizer = amp.initialize(model, optimizer, opt_level='O2', loss_scale=32.0)
优化后代码如下:
optimizer = apex.optimizers.NpuFusedSGD(model.parameters(), lr=args.lr, momentum=args.momentum) #NpuFusedSGD为亲和API函数 model, optimizer = amp.initialize(model, optimizer, opt_level='O2', loss_scale=32.0, combine_grad=True)