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