何为AMP?
全称:Automatic mixed precision自动混合精度。
功能:在神经网络推理过程中,实现针对不同层采用不同的数据精度进行计算,从而实现节省显存和加速训练的目的。
此处提到的不同数据精度包括:32位浮点型torch.FloatTensor()和16位半精度浮点型torch.HalfTensor()。
注意:
1.torch默认精度为32位。
2.amp的使用仅当网络训练在gpu上时,torch.cuda.amp
为何使用AMP?
考虑到32位和16位更有优缺点,想要综合分段采用32位和16位,消除各自的劣势,帮助网络更高效训练。
- 当网络仅采用32位训练时—>优势:32位能很好的将较小的精度数值保存下来完成梯度计算。劣势:batch_size不敢设置过大,对于设备的数据存储能力提出较高要求等。
- 当网络仅采用16位训练时—>优势:存储小,计算快,减少显存占用。劣势:数值范围小,容易出现overflow or underflow。舍入误差,由于误差数值达不到16bit,导致梯度丢失。
何时使用AMP?
amp中的autocast会根据PyTorch框架,进行自动转换,如一般在cond、linear、激活函数等操作中自动转换成16位精度计算。
如何使用AMP?
AMP中包含两个功能模块,分别位autocast和GradScaler。
- torch.cuda.amp.autocast():自动切换32位和16位,但可能会在切换的过程中导致精度的损失,而出现梯度丢失,产生loss NAN问题。
yolox/core/trainer.py中对应的代码展示:
可以看出,它只存在于前向计算中,而未参加反向传播。 - torch.cuda.amp.GradScaler():通过放大loss值来防止梯度消失。
注意:这里面的放大scaler大小是依靠每次迭代时动态估计获得,动态估计指每次epoch时都看一眼有没有出现inf or NAN等错误,若没有出现则尽可能的增大scaler值。
yolox/core/trainer.py中对应的代码展示:
可以看出,它用在反向传播中,但是却没有参与权重梯度的更新(只有一个step())。