前言
使用混合精度训练可以减少显存开销,加快模型的训练速度。更详细的apex
的优缺点可以参考这篇博文:Pytorch自动混合精度(AMP)训练。我们这里介绍的是单GPU
的混合精度训练,多GPU
的训练可以参考下这篇博文英伟达APEX,多GPU分布式训练,同步Batchnorm,自动混合精度训练法宝指南。
一.apex的安装
apex
无法直接通过pip install apex
来安装。下面介绍两种安装方法:
在安装之前请先安装git
工具,这里附上链接:git工具下载,没有别的需求就直接一路next
就行了,安装后记得配置环境变量,自行百度一下即可,很简单。
1.1.官方推荐
这种方法不是很推荐,试了几次都没成功,不知道原因在哪,有知道的伙伴,欢迎评论区留言。
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./2.
1.2.通过git下载文件安装
通过此方法安装成功。这里也需要注意,需要在自己的python环境下面安装。
git clone https://github.com/NVIDIA/apex
cd apex
python setup.py install
二.apex的使用
apex
的使用也比较简单,直接导入apex
包在对模型和参数初始化,以及loss
加入到apex
中即可,下面给出示例:
- 模型和参数部分:
mdoel = Net()
optimizer = AdaXW(params, opt.lr, weight_decay=1e-4)
# 加上这句话即可
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是“欧一”,不是“零一”
- loss部分:
loss = criterion(outputs, labels)
# loss后面加上这两句
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()