Phân tích mã nguồn MMDetection: Cấu trúc RPN Head trong Faster R-CNN

Trong cấu hình mô hình Faster R-CNN sử dụng ResNet-50 và FPN, phần rpn_head đóng vai trò then chốt trong việc đề xuất vùng ứng viên (region proposals). Dưới đây là phân tích chi tiết về lớp RPNHead — thành phần chịu trách nhiệm sinh anchor, dự đoán độ tin cậy và điều chỉnh bounding box.

Cấu hình RPN Head

rpn_head=dict(
    type='RPNHead',
    in_channels=256,
    feat_channels=256,
    anchor_generator=dict(
        type='AnchorGenerator',
        scales=[8],
        ratios=[0.5, 1.0, 2.0],
        strides=[4, 8, 16, 32, 64]),
    bbox_coder=dict(
        type='DeltaXYWHBBoxCoder',
        target_means=[.0, .0, .0, .0],
        target_stds=[1.0, 1.0, 1.0, 1.0]),
    loss_cls=dict(type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0),
    loss_bbox=dict(type='L1Loss', loss_weight=1.0))
  • anchor_generator: Tạo các anchor với tỉ lệ và kích thước đa dạng trên nhiều mức feature map.
  • bbox_coder: Mã hóa/giải mã tọa độ bounding box dưới dạng delta so với anchor.
  • loss_cls & loss_bbox: Hàm mất mát cho phân loại (foreground/background) và hồi quy tọa độ.

Cấu trúc mã nguồn RPNHead

import torch
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import normal_init
from mmcv.ops import batched_nms

from ..builder import HEADS
from .anchor_head import AnchorHead
from .rpn_test_mixin import RPNTestMixin


@HEADS.register_module()
class RPNHead(RPNTestMixin, AnchorHead):

    def __init__(self, in_channels, **kwargs):
        super().__init__(1, in_channels, **kwargs)

    def _build_layers(self):
        self.conv_feat = nn.Conv2d(in_channels=self.in_channels, 
                                   out_channels=self.feat_channels, 
                                   kernel_size=3, padding=1)
        self.conv_score = nn.Conv2d(self.feat_channels, 
                                    self.num_anchors * self.cls_out_channels, 
                                    kernel_size=1)
        self.conv_delta = nn.Conv2d(self.feat_channels, 
                                    self.num_anchors * 4, 
                                    kernel_size=1)

    def init_weights(self):
        for layer in [self.conv_feat, self.conv_score, self.conv_delta]:
            normal_init(layer, std=0.01)

    def forward_per_level(self, feature_map):
        x = F.relu(self.conv_feat(feature_map), inplace=True)
        scores = self.conv_score(x)
        deltas = self.conv_delta(x)
        return scores, deltas

    def compute_loss(self, cls_preds, reg_preds, gt_boxes, meta_info, ignored_boxes=None):
        base_losses = super().loss(cls_preds, reg_preds, gt_boxes, None, meta_info, ignored_boxes)
        return {
            'rpn_cls_loss': base_losses['loss_cls'],
            'rpn_reg_loss': base_losses['loss_bbox']
        }

    def generate_proposals(self, cls_scores, bbox_deltas, anchors_per_level, 
                           image_shape, scale_factors, config, rescale=False):
        cfg = self.test_cfg if config is None else config
        all_scores, all_deltas, all_anchors, level_ids = [], [], [], []

        for level_idx, (score_map, delta_map, anchors) in enumerate(zip(cls_scores, bbox_deltas, anchors_per_level)):
            H, W = score_map.shape[-2:]
            score_map = score_map.permute(1, 2, 0).reshape(-1)
            delta_map = delta_map.permute(1, 2, 0).reshape(-1, 4)
            
            probs = score_map.sigmoid()
            if cfg.nms_pre > 0 and len(probs) > cfg.nms_pre:
                topk_vals, topk_idx = torch.topk(probs, cfg.nms_pre)
                probs = topk_vals
                delta_map = delta_map[topk_idx]
                anchors = anchors[topk_idx]

            all_scores.append(probs)
            all_deltas.append(delta_map)
            all_anchors.append(anchors)
            level_ids.append(probs.new_full((len(probs),), level_idx, dtype=torch.long))

        final_scores = torch.cat(all_scores)
        final_anchors = torch.cat(all_anchors)
        final_deltas = torch.cat(all_deltas)
        level_labels = torch.cat(level_ids)

        proposals = self.bbox_coder.decode(final_anchors, final_deltas, max_shape=image_shape)

        # Lọc box quá nhỏ
        widths = proposals[:, 2] - proposals[:, 0]
        heights = proposals[:, 3] - proposals[:, 1]
        valid_mask = (widths >= cfg.min_bbox_size) & (heights >= cfg.min_bbox_size)
        if not valid_mask.all():
            proposals = proposals[valid_mask]
            final_scores = final_scores[valid_mask]
            level_labels = level_labels[valid_mask]

        # Áp dụng NMS theo cấp độ
        nms_config = dict(type='nms', iou_threshold=cfg.nms_thr)
        final_boxes, _ = batched_nms(proposals, final_scores, level_labels, nms_config)
        return final_boxes[:cfg.nms_post]

Chức năng chính của RPNHead

  1. Khởi tạo: Thiết lập số lớp đầu ra và kênh đặc trưng.
  2. Xây dựng mạng con: Gồm 3 lớp tích chập — trích xuất đặc trưng, phân loại, và hồi quy tọa độ.
  3. Khởi tạo trọng số: Sử dụng phân phối chuẩn với độ lệch chuẩn 0.01.
  4. Forward từng mức: Xử lý từng feature map từ FPN để sinh điểm và delta box.
  5. Tính toán loss: Kế thừa từ lớp cha, trả về loss phân loại và hồi quy.
  6. Sinh đề xuất cuối: Giải mã tọa độ, lọc box nhỏ, và áp dụng NMS theo mức độ feature map.

Thẻ: mmdetection faster-rcnn rpn-head PyTorch object-detection

Đăng vào ngày 23 tháng 5 lúc 00:50