利用NVIDIA的apex实现在Pytorch中混合精度计算。

利用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基本等于最开始实验的全半精度的运算。值得一提的是，不论在优化过程中，模型是否采用半精度，保存下来的模型均为单精度模型，能够保证模型在其他应用中的正常使用。