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

it2022-05-05  91

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

 

 

   

 

 


最新回复(0)