apex混合精度训练教程

前言

  使用混合精度训练可以减少显存开销，加快模型的训练速度。更详细的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()