Pytroch使用混合精度进行训练

原创已于 2022-08-16 17:33:01 修改 · 574 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#linux #pytorch #运维

于 2022-07-10 19:28:01 首次发布

Pytroch基础专栏收录该内容

18 篇文章

订阅专栏

本文分享了在使用PyTorch进行GPU混合精度训练时遇到的问题，如Apex库的安装困扰及GradScaler的使用调整。通过详细步骤和解决方法，帮助读者顺利实现高效训练。

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

近期使用Pytroch在GPU上进行混合精度训练出现一些问题，通过不断努力最终解决了，在此进行分享。

一、Apex安装

网上在LInux系统上关于apex库的安装主要为步骤如下：

$ git clone https://github.com/NVIDIA/apex
$ cd apex
$ python3 setup.py install

但是通过上述步骤操作后，我得到的apex库引用出现了一些问题。
安装apex库出现问题的解决方法
这个网站介绍了安装apex库出现问题的解决方法。最后我通过
apex库的安装
这个网址提供的方法成功安装了apex库：
For performance and full functionality, we recommend installing Apex with CUDA and C++ extensions via

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

Apex also supports a Python-only build via（通过python安装，把上述代码第三句替换一下）

pip install -v --disable-pip-version-check --no-cache-dir ./

二、apex的使用

关于apex的使用根据pytroch的版本不同有着不同的使用方法。主要参考以下文章：
混合精度的使用
其中使用apex会出现一些小问题，我主要出现Gradscaler()无法调用，最终将scaler = GradScaler()替换为scale=torch.cuda.amp.GradScaler()
解决了问题。
举一个torch版本小于等于1.5使用O2+动态loss scale的例子。

from apex import amp
try:
    from apex import amp
    APEX_AVAILABLE = True
except ModuleNotFoundError:
    APEX_AVAILABLE = False

# 这里的loss和optimizer取自定义的损失函数和优化器
def amp_backward(loss, optimizer):
    if APEX_AVAILABLE:
        with amp.scale_loss(loss, optimizer) as scaled_loss:
            scaled_loss.backward()
    else:
        loss.backward()

# 这里的net和optimizer取自定义的网络和优化器
if APEX_AVAILABLE:
        net, optimizer = amp.initialize(
            net, optimizer, opt_level="O2", keep_batchnorm_fp32=True,
            loss_scale="dynamic"
        )    # 使用动态混合精度

# 将loss.backward()替换为下面这段代码
amp_backward(loss, optimizer)