利用NVIDIA的apex实现在Pytorch中混合精度计算。
参考链接:https://devblogs.nvidia.com/apex-pytorch-easy-mixed-precision-training/
https://nvidia.github.io/apex/amp.html
https://github.com/NVIDIA/apex
1. 在linux上安装:
(1) 建议使用CUDA和C ++扩展安装Apex
$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
(2) 也可以仅安装支持python版本的(需要Pytorch 0.4)
$ pip install -v --no-cache-dir ./
2.简单的使用
#首先import amp
from apex import amp
......
#apm可以选择不同的opt_level
model = model.cuda()
model, optimizer = amp.initialize(model, optimizer, opt_level="O2")
......
#loss.backward()
optimizer.zero_grad() with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward()
3.opt_level简单解释
opt_level有四种可选值O0,O1,O2,O3
概括起来:00相当于原始的单精度训练。01在大部分计算时采用半精度,但是所有的模型参数依然保持单精度,对于少数单精度较好的计算(如softmax)依然保持单精度。02相比于01,将模型参数也变为半精度。03基本等于最开始实验的全半精度的运算。值得一提的是,不论在优化过程中,模型是否采用半精度,保存下来的模型均为单精度模型,能够保证模型在其他应用中的正常使用。