深度学习训练时冻结部分参数的方法

12 篇文章 0 订阅

问题描述

在使用mmdetection代码库时需要冻结部分网络参数,只训练一部分的网络。这里提供一种简单且不容易出现bug的方法,不仅仅适用于mmdetection代码库,也可以使用在其他的代码库里面,不过需要一定的改动。

解决方案

mmdetection这个库里面优化器的初始化的位置在:mmdet/apis/train.py,如果使用的是其他代码库的话,需要找到对应的optimizer的初始化的位置。在mmdetection中,原始的optimizer的定义为:

# build optimizer
auto_scale_lr(cfg, distributed, logger)
optimizer = build_optimizer(model, cfg.optimizer)

将这里原本optimizer定义的方式去掉(注释掉),改为新的定义方式。

首先,确定自己需要使用的优化器,一般从config文件里就可以确认。比如mmdetection里面,configs/pascal_voc/faster_rcnn_r50_fpn_1x_voc0712.py 中就有:

optimizer = dict(type='SGD', lr=0.001, momentum=0.9, weight_decay=0.0001)

我们根据config文件确认自己需要使用的是SGD优化器之后,直接通过pytorch生成优化器。通过这种方式定义优化器时,需要指定要优化的参数。这样,我们就可以冻结部分参数,对另一部分参数进行优化。具体来说,可以先打印模型中的参数,确认自己要优化的参数的name之后,将其加入到待优化的参数列表里面,如下面的代码所示。

parameters = []
for name, p in model.named_parameters():
    print(name)
    if "retina_cls" in name:
        parameters.append(p)

optimizer = torch.optim.SGD(
    parameters, lr=cfg.optimizer['lr'], momentum=cfg.optimizer['momentum'],
    weight_decay=cfg.optimizer['weight_decay']
)

完成上面的操作之后,除了parameters里面的参数其他都不会被优化、不会变了。有的同学可能不太放心,那么在实际的model里可以打印参数的具体数值、回传梯度、requires_grad等信息来确认。比如,我使用的是faster RCNN网络,找到forward_train函数。faster RCNN的forward_train函数是定义在其父类TwoStageDetector里(位置在mmdet/models/detectors/two_stage.py)。在loss回传前加入如下代码,打印参数的name、requires_grad以及具体参数数值。

for name, p in self.roi_head.named_parameters():
    print(name)
    print(p.requires_grad)
    print(p)
 
losses.update(roi_losses)

打印之后可以发现,虽然所有的参数requires_grad都为True,都具有回传梯度,但是除了我们在parameters里面指定的参数,其余的都没有更新(反复运行几次看数值有没有变化即可)。值得注意的是,我们不需要担心其他参数会随着梯度而更新,所以额外再设定其他的参数requires_grad为false了,这样会导致程序出现bug。实际上,其余的参数虽然有回传梯度,但是因为不在optimizer优化范围内,所以不会更新

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

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值