如何在骨干网络已预训练的torchvision Faster R-CNN中仅训练RPN模块

目前torchvision的FasterRCNN构造函数没有提供直接通过传参实现仅训练RPN的配置，你需要手动冻结不需要训练的模块参数即可，不需要大改训练逻辑，操作步骤如下：

操作步骤 #

• 第一步：冻结不需要训练的模块参数
  你不需要训练的是RoI Heads模块（对应后续的目标分类、检测框二次回归部分），如果你的骨干网络已经完成预训练不需要微调，也可以同步冻结骨干参数，代码如下：

  # 情况1：骨干网络不需要微调，直接冻结
  for param in fastercnn_model.backbone.parameters():
      param.requires_grad = False
  
  # 冻结后续分类器和框回归模块
  for param in fastercnn_model.roi_heads.parameters():
      param.requires_grad = False
  
  # 确保RPN参数处于可训练状态
  for param in fastercnn_model.rpn.parameters():
      param.requires_grad = True
  

  如果你需要同步微调骨干网络，删除冻结backbone的代码段即可。
• 第二步：构造优化器时仅传入可训练参数
  这样可以避免无用的梯度计算，降低显存占用：

  import torch.optim as optim
  
  optimizer = optim.SGD(
      params=[p for p in fastercnn_model.parameters() if p.requires_grad],
      lr=0.005,
      momentum=0.9,
      weight_decay=5e-4
  )
  

注意事项 #

训练时你可以直接调用torchvision Faster R-CNN的原生前向逻辑，即loss_dict = fastercnn_model(images, targets)，框架会自动计算RPN的分类损失和回归损失，RoI Heads部分的参数已经被冻结，即使计算了对应损失也不会更新参数。如果追求极致的训练速度，你也可以修改模型前向逻辑跳过RoI Heads的计算，常规场景下直接冻结参数即可满足需求。

内容的提问来源于stack exchange，提问作者Anh Minh Tran