如何将Tuned Faster R-CNN模型最大检测目标数提升至3000？

问题背景 #

当前使用PyTorch的Faster RCNN模型处理图片时，单张图存在300+个目标，但模型仅能检测出100个，需要将最大检测目标数提升至3000。

用于准备待优化模型的代码：

# load Faster RCNN pre-trained model
Faster_RCNN_tuned_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT")

# Get the number of input features 
in_features = Faster_RCNN_tuned_model.roi_heads.box_predictor.cls_score.in_features
    
# Define a new head for the detector with 2 classes (cell or fone)
Faster_RCNN_tuned_model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2) 
Faster_RCNN_tuned_model = Faster_RCNN_tuned_model.to(DEVICE)

加载微调后模型的代码：

# Create Faster RCNN default model
best_Faster_RCNN_tuned_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT")
# Get the number of input features 
in_features = best_Faster_RCNN_tuned_model.roi_heads.box_predictor.cls_score.in_features
    
# Define a new head for the detector with 2 classes (cell or fone)
best_Faster_RCNN_tuned_model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2) 
# Load and setup parameters from saved best model
checkpoint = torch.load('C:\\temp\\datasets\\mediag\\models\\best_model.pth', map_location=DEVICE)
best_Faster_RCNN_tuned_model.load_state_dict(checkpoint['model_state_dict'])
best_Faster_RCNN_tuned_model = best_Faster_RCNN_tuned_model.to(DEVICE).eval()

解决方案 #

需要手动调整Faster RCNN的两个核心参数，来提升最大检测目标数：

1. 调整RPN阶段候选框数量 #

RPN（区域建议网络）在NMS（非极大值抑制）后默认保留的候选框数量有限，需修改训练/推理阶段的候选框上限值，设置为足够覆盖目标数量的数值（如3000）。

2. 调整最终检测框数量 #

RoI Heads阶段默认限制单张图最多输出100个检测框，需修改该上限值为目标数量（如3000）。

修改后的代码示例 #

初始化待优化模型时配置参数

# 加载预训练模型并配置检测数量参数
model_config = {
    "weights": "DEFAULT",
    "rpn_post_nms_top_n_train": 3000,  # 训练阶段RPN保留候选框数
    "rpn_post_nms_top_n_test": 3000,   # 推理阶段RPN保留候选框数
    "detections_per_img": 3000          # 单图最大检测框数
}
Faster_RCNN_tuned_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(**model_config)

# 替换自定义预测头
in_features = Faster_RCNN_tuned_model.roi_heads.box_predictor.cls_score.in_features
Faster_RCNN_tuned_model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2) 
Faster_RCNN_tuned_model = Faster_RCNN_tuned_model.to(DEVICE)

加载已微调模型后调整参数

如果是加载已训练完成的模型，直接修改模型属性即可：

# 创建模型并替换预测头
best_Faster_RCNN_tuned_model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT")
in_features = best_Faster_RCNN_tuned_model.roi_heads.box_predictor.cls_score.in_features
best_Faster_RCNN_tuned_model.roi_heads.box_predictor = FastRCNNPredictor(in_features, 2) 

# 加载预训练权重
checkpoint = torch.load('C:\\temp\\datasets\\mediag\\models\\best_model.pth', map_location=DEVICE)
best_Faster_RCNN_tuned_model.load_state_dict(checkpoint['model_state_dict'])

# 调整检测数量参数
best_Faster_RCNN_tuned_model.rpn.post_nms_top_n_test = 3000
best_Faster_RCNN_tuned_model.roi_heads.detections_per_img = 3000

best_Faster_RCNN_tuned_model = best_Faster_RCNN_tuned_model.to(DEVICE).eval()

注意事项 #

• 提升候选框和检测框数量会增加推理时间，需根据硬件性能平衡检测数量与速度。
• 如果训练阶段也需要处理大量目标，必须同步调整rpn_post_nms_top_n_train参数，避免训练与推理阶段参数不一致。

内容的提问来源于stack exchange，提问作者Itai Kogtev