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错误：https://blog.csdn.net/u013823233/article/details/101209718