Implementing Backbone and Feature Pyramid Network#

首先构建 backbone（骨干网络） 和 FPN（特征金字塔）如图示中蓝色部分和绿色部分

具体使用 RegNetX-400MF 从预训练的ImageNet权重初始化这个backbone，提取中间特征(c3, c4, c5)，这些特征(c3, c4, c5)的高度和宽度分别是输入图像的 $1/8^{th}, 1/16^{th}, 1/32^{th}$

这些值(8, 16, 32)被称为这些特征的 “stride”

For dummy input images with shape: (2, 3, 224, 224)
Shape of c3 features: torch.Size([2, 64, 28, 28])
Shape of c4 features: torch.Size([2, 160, 14, 14])
Shape of c5 features: torch.Size([2, 400, 7, 7])

实现FPN提取三个特征 class DetectorBackboneWithFPN(nn.Module)

具体思路：（为了方便训练这里仅使用提取出的3个特征）

通过backbone提取出的3个特征经过kernel_size = 1, stride=1的统一通道数，然后在从特征图大小最小的c5开始进行上采样到和c4特征图大小相同，相加融合之后得到m4, 对m4再进行上采样到和c3大小相同，相加融合得到m3

在对m3, m4, m5进行kernel_size = 3, stride=1, padding=1的卷积(特征图大小不变)，输出p3, p4, p5

class DetectorBackboneWithFPN(nn.Module):

    def __init__(self, out_channels: int):
        self.fpn_params = nn.ModuleDict({
                "m3": nn.Conv2d(in_channels=dummy_out_shapes[0][1][1], out_channels=out_channels, kernel_size=1, stride=1),
                "m4": nn.Conv2d(in_channels=dummy_out_shapes[1][1][1], out_channels=out_channels, kernel_size=1, stride=1),
                "m5": nn.Conv2d(in_channels=dummy_out_shapes[2][1][1], out_channels=out_channels, kernel_size=1, stride=1),
                
                "p3": nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1),
                "p4": nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1),
                "p5": nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1),
            })

    def forward(self, images: torch.Tensor):

        m3 = self.fpn_params["m3"](backbone_feats["c3"])
        m4 = self.fpn_params["m4"](backbone_feats["c4"])
        m5 = self.fpn_params["m5"](backbone_feats["c5"])
        
        m5_upsample = F.interpolate(m5, size=(m4.shape[2], m4.shape[3]), mode="nearest")
        m4 = m4 + m5_upsample
        
        m4_upsample = F.interpolate(m4, size=(m3.shape[2], m3.shape[3]), mode="nearest")
        m3 = m3 + m4_upsample
        
        fpn_feats["p3"] = self.fpn_params["p3"](m3)
        fpn_feats["p4"] = self.fpn_params["p4"](m4)
        fpn_feats["p5"] = self.fpn_params["p5"](m5)

        return fpn_feats

Implementing FCOS prediction network (head)#

实现头部预测块（原论文中的橙色块），其内部含有4个卷积层(对应代码中stem_channels的长度为4)，输出(H, W, 256)大小的特征（不同FPN层级的H, W不同）然后每个卷积层使用一个卷积层进行最终预测

在预期输出中，classification logits 有 NUM_CLASSES 个通道，box regression deltas 有 4 个输出通道，centerness 有 1 个输出通道。 所有输出的高度和宽度被展平为一个维度，从而得到 (B, H * W, C) 格式 - 这种格式更方便计算损失.

class FCOSPredictionNetwork(nn.Module):
    """
    FCOS prediction network that accepts FPN feature maps from different levels
    and makes three predictions at every location: bounding boxes, class ID and
    centerness. This module contains a "stem" of convolution layers, along with
    one final layer per prediction. For a visual depiction, see Figure 2 (right
    side) in FCOS paper: https://arxiv.org/abs/1904.01355

    We will use feature maps from FPN levels (P3, P4, P5) and exclude (P6, P7).
    """

    def __init__(
        self, num_classes: int, in_channels: int, stem_channels: List[int]
    ):
        """
        Args:
            num_classes: Number of object classes for classification.
            in_channels: Number of channels in input feature maps. This value
                is same as the output channels of FPN, since the head directly
                operates on them.
            stem_channels: List of integers giving the number of output channels
                in each convolution layer of stem layers.
        """
        super().__init__()

        # Fill these.
        stem_cls = []
        stem_box = []
        # Replace "pass" statement with your code
        prev_out = in_channels
        for output_channel in stem_channels:
            conv_cls = nn.Conv2d(in_channels=prev_out, out_channels=output_channel, kernel_size=3, stride=1, padding=1)
            nn.init.normal_(conv_cls.weight.data, mean=0, std=0.01)
            nn.init.constant_(conv_cls.bias.data, 0)
            stem_cls.append(conv_cls)
            stem_cls.append(nn.ReLU())
            
            conv_box = nn.Conv2d(in_channels=prev_out, out_channels=output_channel, kernel_size=3, stride=1, padding=1)
            nn.init.normal_(conv_box.weight.data, mean=0, std=0.01)
            nn.init.constant_(conv_box.bias.data, 0)
            stem_box.append(conv_box)
            stem_box.append(nn.ReLU())
            
            prev_out = output_channel

        self.stem_cls = nn.Sequential(*stem_cls)
        self.stem_box = nn.Sequential(*stem_box)


        self.pred_cls = None  # Class prediction conv
        self.pred_box = None  # Box regression conv
        self.pred_ctr = None  # Centerness conv

        # Replace "pass" statement with your code
        self.pred_cls = nn.Conv2d(in_channels=prev_out, out_channels=num_classes, kernel_size=3, stride=1, padding=1)
        self.pred_box = nn.Conv2d(in_channels=prev_out, out_channels=4, kernel_size=3, stride=1, padding=1)
        self.pred_ctr = nn.Conv2d(in_channels=prev_out, out_channels=1, kernel_size=3, stride=1, padding=1)
        
        for layer in [self.pred_cls, self.pred_box, self.pred_ctr]:
            nn.init.normal_(layer.weight.data, mean=0, std=0.01)
            nn.init.constant_(layer.bias.data, 0)

        # OVERRIDE: Use a negative bias in `pred_cls` to improve training
        # stability. Without this, the training will most likely diverge.
        torch.nn.init.constant_(self.pred_cls.bias, -math.log(99))

    def forward(self, feats_per_fpn_level: TensorDict) -> List[TensorDict]:

        class_logits = {}
        boxreg_deltas = {}
        centerness_logits = {}

        # Replace "pass" statement with your code
        for key, value in feats_per_fpn_level.items():
            class_logits[key] = self.pred_cls(self.stem_cls(value))     # (N, C, H, W)
            N, C, H, W = class_logits[key].shape
            class_logits[key] = class_logits[key].reshape(N, C, -1)
            class_logits[key] = class_logits[key].permute(0, 2, 1)

            boxreg_deltas[key] = self.pred_box(self.stem_box(value))    # (N, 4, H, W)
            boxreg_deltas[key] = boxreg_deltas[key].reshape(N, 4, -1)
            boxreg_deltas[key] = boxreg_deltas[key].permute(0, 2, 1)
            
            centerness_logits[key] = self.pred_ctr(self.stem_box(value))    # (N, 1, H, W)
            centerness_logits[key] = centerness_logits[key].reshape(N, 1, -1)
            centerness_logits[key] = centerness_logits[key].permute(0, 2, 1)

        return [class_logits, boxreg_deltas, centerness_logits]

Assigning a GT target to every model prediction#

首先根据每层特征图生成location(特征图中心点)，由于特征图是在原图的基础上缩小了，需要将这些中心点的坐标放大回原图的坐标系中去

def get_fpn_location_coords(
    shape_per_fpn_level: Dict[str, Tuple],
    strides_per_fpn_level: Dict[str, int],
    dtype: torch.dtype = torch.float32,
    device: str = "cpu",
) -> Dict[str, torch.Tensor]:
    
    location_coords = {
        level_name: None for level_name, _ in shape_per_fpn_level.items()
    }

    for level_name, feat_shape in shape_per_fpn_level.items():
        level_stride = strides_per_fpn_level[level_name]

        N, C, H, W = feat_shape
        
        rows = torch.arange(start=0, end=H, dtype=dtype, device=device)
        cols = torch.arange(start=0, end=W, dtype=dtype, device=device)
        
        y_grid, x_grid = torch.meshgrid(rows, cols, indexing="ij")
        
        xc = level_stride * (x_grid + 0.5)
        yc = level_stride * (y_grid + 0.5)
        
        stack_location_coords = torch.stack([xc, yc], dim=-1)
        location_coords[level_name] = stack_location_coords.reshape(-1, 2)

    return location_coords

NOTE：“feature map location” 和 “feature center”表示相同的意思，就是对特征图划分成网格后取中心点(图中的黄点)

Matching feature map locations with GT boxes#

在将特征图中心点映射到原图坐标系中去之后，需要进行确定中心点匹配哪些真实标注的GT框，参考论文使用两种匹配规则：

Rule 1: 取中心点到4个GT框的边界距离的最小值，如果存在小于0，则不匹配

Rule 2: 根据FPN层级，设定多匹配规则：

p4, p5的下限为 stride * 4, p3的下限为0

p3, p4的上限为 stride * 8，p5为无穷不设置上限

def fcos_match_locations_to_gt(
    locations_per_fpn_level: TensorDict,
    strides_per_fpn_level: Dict[str, int],
    gt_boxes: torch.Tensor,
) -> TensorDict:
    
    matched_gt_boxes = {
        level_name: None for level_name in locations_per_fpn_level.keys()
    }

    # Do this matching individually per FPN level.
    for level_name, centers in locations_per_fpn_level.items():

        # Get stride for this FPN level.
        stride = strides_per_fpn_level[level_name]

        x, y = centers.unsqueeze(dim=2).unbind(dim=1)
        x0, y0, x1, y1 = gt_boxes[:, :4].unsqueeze(dim=0).unbind(dim=2)
        pairwise_dist = torch.stack([x - x0, y - y0, x1 - x, y1 - y], dim=2)

        # Pairwise distance between every feature center and GT box edges:
        # shape: (num_gt_boxes, num_centers_this_level, 4)
        pairwise_dist = pairwise_dist.permute(1, 0, 2)

        # The original FCOS anchor matching rule: anchor point must be inside GT.
        match_matrix = pairwise_dist.min(dim=2).values > 0

        # Multilevel anchor matching in FCOS: each anchor is only responsible
        # for certain scale range.
        # Decide upper and lower bounds of limiting targets.
        pairwise_dist = pairwise_dist.max(dim=2).values

        lower_bound = stride * 4 if level_name != "p3" else 0
        upper_bound = stride * 8 if level_name != "p5" else float("inf")
        match_matrix &= (pairwise_dist > lower_bound) & (
            pairwise_dist < upper_bound
        )

        # Match the GT box with minimum area, if there are multiple GT matches.
        gt_areas = (gt_boxes[:, 2] - gt_boxes[:, 0]) * (
            gt_boxes[:, 3] - gt_boxes[:, 1]
        )

        # Get matches and their labels using match quality matrix.
        match_matrix = match_matrix.to(torch.float32)
        match_matrix *= 1e8 - gt_areas[:, None]

        # Find matched ground-truth instance per anchor (un-matched = -1).
        match_quality, matched_idxs = match_matrix.max(dim=0)
        matched_idxs[match_quality < 1e-5] = -1

        # Anchors with label 0 are treated as background.
        matched_boxes_this_level = gt_boxes[matched_idxs.clip(min=0)]
        matched_boxes_this_level[matched_idxs < 0, :] = -1

        matched_gt_boxes[level_name] = matched_boxes_this_level

    return matched_gt_boxes