apex混合精度训练教程

I松风水月

已于 2023-04-04 09:09:48 修改

阅读量1k

点赞数 1

CC 4.0 BY-SA版权

于 2023-04-03 11:25:44 首次发布

本文链接：https://blog.youkuaiyun.com/qq_38683460/article/details/129925312

pytorch 专栏收录该内容

14 篇文章

订阅专栏

文章介绍了如何在Pytorch中利用NVIDIA的apex库进行混合精度训练，以减少显存使用并加速模型训练。apex的安装包括官方推荐但可能不成功的gitclone方法和一种成功的通过pythonsetup.pyinstall的方式。在使用apex时，需对模型和优化器进行初始化，并在损失函数中加入特定的scale_loss和backward操作。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

前言

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