mmdetection添加apex训练模型

1. 首先确保不加apex前模型能训练起来,精度正常

2. 在环境上安装好apex:

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" ./

3. 修改模型初始化过程,加上apex初始化

修改文件:mmdet/tools/train.py

导入apex包

修改amp初始化:

文件修改diff如下:

4. 修改梯度反向写法

首先弄清楚mmdetection反向的过程,这部分是调用的mmcv实现

runner中实现了训练的过程,以最基本的runner为例,文件:mmcv/runner/epoch_based_runner.py

在self.run_iter之后,调用hook函数来拉起训练后续操作(比如loss.backward和optimizer.step过程) ,具体的函数有:

可以看到after_train_iter后调用到OptimizerHook中的after_train_iter函数,进一步跟入到文件:mmcv/runner/hooks/optimizer.py

看到了熟悉的loss.backward过程,将图中的runner.outputs['loss'].backward()修改为:

with amp.scale_loss(runner.outputs['loss'], runner.optimizer) as scaled_loss:
    scaled_loss.backward()

至此修改完成

重跑训练,验证精度是否OK即可。

参考:

apex使用:https://github.com/NVIDIA/apex

docker运行shm错误:ERROR: Unexpected bus error encountered in worker. This might be caused by insufficient shared memo_u013823233的专栏-CSDN博客

おすすめ

転載: blog.csdn.net/dou3516/article/details/121283989