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博客

  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值