SOLOv2（NeurIPS 2020）论文与代码解读

paper：SOLOv2: Dynamic and Fast Instance Segmentation

official implementation：AdelaiDet/configs/SOLOv2/README.md at master · aim-uofa/AdelaiDet · GitHub

third-party implementation：mmdetection/mmdet/models/dense_heads/solov2_head.py at main · open-mmlab/mmdetection · GitHub

存在的问题

作者指出在SOLO中，有三个主要的瓶颈限制了它的性能：

1. 低效的mask表示和学习
2. 没有足够大的分辨率来进行更精细的mask预测
3. mask NMS速度慢

创新点

针对上述三个问题，本文提出了对应的解决方法，并将新的框架称为SOLO v2：

• 动态实例分割：SOLOv2将掩码学习过程分为卷积核学习和特征学习。通过动态预测卷积核，掩码生成变得更灵活和高效。
• 掩码核（Mask Kernel）：在每个特征金字塔级别上预测掩码核，采用坐标卷积（CoordConv）添加空间功能。输出的卷积核权重根据位置进行预测。
• 掩码特征（Mask Feature）：统一预测所有FPN级别的掩码特征表示，通过特征金字塔融合构建高分辨率的掩码特征表示。
• Matrix NMS：引入矩阵NMS，通过并行矩阵操作一次性进行NMS，减少重复预测，提高速度和精度。

方法介绍

在SOLO v1中，输入图片被划分为 (S\times S) 个网格，如果一个目标的中心点落入某个网格中，该网格负责该目标的binary mask的预测。因此模型一共输出 (S^2) 个mask，表示为 (M\in \mathbb{R}^{H\times W\times S^2})。第 (k) 个通道负责分割位置 ((i,j)) 的实例，其中 (k=i\cdot S+j)，如图2(a)所示。

这种范式可以优雅地生成实例分割结果，但有三个瓶颈限制了其性能：a）低效的掩码表示和学习。预测具有 (S^2) 个通道的输出张量 (M) 需要大量的内存和计算。此外，由于不同FPN level的 (S) 不同，每个level的最后一层是单独学习的不共享，这就导致了训练效率低下。b）不准确的mask预测。更精细的预测需要高分辨率的mask来处理物体边缘上的细节。但大分辨率大大增加了计算成本。c）maks NMS速度慢。和box NMS相比，mask NMS需要更多的处理时间导致开销更大。

Dynamic Instance Segmentation

在SOLO v1中，为了生成对应 (S\times S) 个网格的具有 (S^2) 个通道的instance mask，最后一层以FPN一个level的特征 (F\in \mathbb{R}^{H\times W\times E}) 为输入，并通过一个卷积层得到 (S^2) 通道的输出

其中 (G_{i,j}\in \mathbb{R}^{1\times 1\times E}) 是卷积核，(M_{i,j}\in \mathbb{R}^{H\times W}) 是最终的mask只包含一个instance其中心在位置 ((i,j))。

换句话说我们需要两个输入 (F) 和 (G) 来生成最终的掩码 (M)。SOLO输出整个 (M) 用于训练和推理，(M) 非常大，直接预测 (M) 对于内存和计算效率都是低下的。在大多数情况下，目标在图像中是稀疏的，因此 (M) 冗余的，因为在一次推理中 (S^2) 个kernel中只有一小部分起作用。

因此本文提出分别学习 (F) 和 (G)，这样就可以从预测的 (S^2) 个kernel中选择有效的然后动态地执行卷积。

Mask Kernel G

给定backbone和FPN，我们在每个pyramid level预测mask kernel。首先将输入特征 (F_I\in\mathbb{R}^{H_I\times W_I\times C}) resize成shape (S\times S\times C)，然后4个卷积和最后一个 (3\times 3\times D) 的卷积用来生成kernel (G)。其中和SOLO一样通过CoordConv向输入特征中添加坐标信息。

对于每个网格，kernel branch预测 (D) 维的输出作为卷积核参数 。(D) 是参数数量，对于输入通道数为 (E) 的1x1卷积 (D) 等于 (E)，对于3x3卷积 (D) 等于 (9E)。

Mask Feature F

由于mask feature和mask kernel解耦了是分别预测的，有两种方法构建mask feature。一是把它和kernel branch一起放到head部分，这意味着我们对每个FPN level分别预测mask feature。另一种是对所有的FPN level预测一个统一的mask representation，作者通过实验比较了两者的效果，最终决定采用后者。

为了学习一个统一的高分辨率的mask特征表示，作者应用了特征金字塔融合。在若干3x3卷积、group norm、ReLU和2x的上采样后，P2到P5的FPN特征融合成一个单独的1/4大小的特征。在elment-wise求和之后最后一层由一个1x1卷积、group norm和ReLU组成。

Matrix NMS

Matrix NMS的具体介绍可参考Fast NMS和Matrix NMS解读-CSDN博客。

代码解析

这里以mmdetection中的实现为例讲解一下代码。输入shape=(1, 3, 736, 1344)，其中batch_size=1。然后进入solov2_head.py中的forward函数，这里输入x就是经过backbone和neck后的输出，是一个列表，包含5个FPN level的输出特征。代码如下

    def forward(self, x):
        """Forward features from the upstream network.

        Args:
            x (tuple[Tensor]): Features from the upstream network, each is
                a 4D-tensor.

        Returns:
            tuple: A tuple of classification scores, mask prediction,
            and mask features.

                - mlvl_kernel_preds (list[Tensor]): Multi-level dynamic kernel
                  prediction. The kernel is used to generate instance
                  segmentation masks by dynamic convolution. Each element in
                  the list has shape
                  (batch_size, kernel_out_channels, num_grids, num_grids).
                - mlvl_cls_preds (list[Tensor]): Multi-level scores. Each
                  element in the list has shape
                  (batch_size, num_classes, num_grids, num_grids).
                - mask_feats (Tensor): Unified mask feature map used to
                  generate instance segmentation masks by dynamic convolution.
                  Has shape (batch_size, mask_out_channels, h, w).
        """
        assert len(x) == self.num_levels
        # [(1,256,184,336),
        #  (1,256,92,168),
        #  (1,256,46,84),
        #  (1,256,23,42),
        #  (1,256,12,21)]
        mask_feats = self.mask_feature_head(x)  # (1,256,184,336)
        ins_kernel_feats = self.resize_feats(x)
        # [(1,256,92,168),
        #  (1,256,92,168),
        #  (1,256,46,84),
        #  (1,256,23,42),
        #  (1,256,23,42)]
        mlvl_kernel_preds = []
        mlvl_cls_preds = []
        for i in range(self.num_levels):
            ins_kernel_feat = ins_kernel_feats[i]
            # ins branch
            # concat coord
            coord_feat = generate_coordinate(ins_kernel_feat.size(),
                                             ins_kernel_feat.device)
            ins_kernel_feat = torch.cat([ins_kernel_feat, coord_feat], 1)  # (1,256,92,168), (1,2,92,168) -> (1,258,92,168)

            # kernel branch
            kernel_feat = ins_kernel_feat  # (1,258,92,168)
            kernel_feat = F.interpolate(
                kernel_feat,
                size=self.num_grids[i],
                mode='bilinear',
                align_corners=False)  # (1,258,40,40)

            cate_feat = kernel_feat[:, :-2, :, :]  # (1,256,40,40)

            kernel_feat = kernel_feat.contiguous()
            for i, kernel_conv in enumerate(self.kernel_convs):  # 所有level共享kernel_convs
                kernel_feat = kernel_conv(kernel_feat)  # (1,512,40,40)

            kernel_pred = self.conv_kernel(kernel_feat)  # (1,256,40,40)

            # cate branch
            cate_feat = cate_feat.contiguous()
            for i, cls_conv in enumerate(self.cls_convs):
                cate_feat = cls_conv(cate_feat)  # (1,512,40,40)
            cate_pred = self.conv_cls(cate_feat)  # (1,1,40,40)

            mlvl_kernel_preds.append(kernel_pred)
            mlvl_cls_preds.append(cate_pred)

        # [(1,256,40,40),
        #  (1,256,36,36),
        #  (1,256,24,24),
        #  (1,256,16,16),
        #  (1,256,12,12)]
        ##########
        # [(1,1,40,40),
        #  (1,1,36,36),
        #  (1,1,24,24),
        #  (1,1,16,16),
        #  (1,1,12,12)]
        ##########
        # (1,256,184,336)
        return mlvl_kernel_preds, mlvl_cls_preds, mask_feats

其中

self.mask_feature_head

就是上面提到的对不同的PFN level特征进行融合得到一个统一特征，这里不贴原始代码了，直接print出module的结构如下。包含0-3共4个module，即不对最小的FPN特征处理。可以看到对于最大的特征就是一个简单的conv-gn-relu，对于1/2大的特征是一个conv-gn-relu再加一个2x上采样然后与第一层的特征进行element-wise相加，对于1/4大的特征则进行两次2x上采样，每个上采样前都是conv-gn-relu，而对于最后一个1/8大的特征则是经过三次2x上采样，注意最后一个module第一层卷积的输入通道数为258是因为通过coordconv加上了xy坐标信息。将4个不同大小的FPN level的特征通过上采样得到相同大小的特征后相加求和，最后通过一个1x1 conv得到输出。

MaskFeatModule(
  (convs_all_levels): ModuleList(
    (0): Sequential(
      (conv0): ConvModule(
        (conv): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (gn): GroupNorm(32, 128, eps=1e-05, affine=True)
        (activate): ReLU()
      )
    )
    (1): Sequential(
      (conv0): ConvModule(
        (conv): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (gn): GroupNorm(32, 128, eps=1e-05, affine=True)
        (activate): ReLU()
      )
      (upsample0): Upsample(scale_factor=2.0, mode=bilinear)
    )
    (2): Sequential(
      (conv0): ConvModule(
        (conv): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (gn): GroupNorm(32, 128, eps=1e-05, affine=True)
        (activate): ReLU()
      )
      (upsample0): Upsample(scale_factor=2.0, mode=bilinear)
      (conv1): ConvModule(
        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (gn): GroupNorm(32, 128, eps=1e-05, affine=True)
        (activate): ReLU()
      )
      (upsample1): Upsample(scale_factor=2.0, mode=bilinear)
    )
    (3): Sequential(
      (conv0): ConvModule(
        (conv): Conv2d(258, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (gn): GroupNorm(32, 128, eps=1e-05, affine=True)
        (activate): ReLU()
      )
      (upsample0): Upsample(scale_factor=2.0, mode=bilinear)
      (conv1): ConvModule(
        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (gn): GroupNorm(32, 128, eps=1e-05, affine=True)
        (activate): ReLU()
      )
      (upsample1): Upsample(scale_factor=2.0, mode=bilinear)
      (conv2): ConvModule(
        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (gn): GroupNorm(32, 128, eps=1e-05, affine=True)
        (activate): ReLU()
      )
      (upsample2): Upsample(scale_factor=2.0, mode=bilinear)
    )
  )
  (conv_pred): ConvModule(
    (conv): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (gn): GroupNorm(32, 256, eps=1e-05, affine=True)
    (activate): ReLU(inplace=True)
  )
)

再看回forward函数，接下来

self.resize_feats

对FPN最大的特征下采样2x对最小的特征上采样2x，和SOLOv1一样。然后进入kernel branch，这里的输入是每个level的特征，且都concat了归一化坐标，然后通过插值F.interpolate resize到 (S^2) 大小。然后

self.kernel_convs

是四层卷积，最后通过

self.conv_kernel

得到卷积核的参数，这里self.kernel_convs和self.conv_kernel对不同level的特征都是共享的，self.conv_kernel的输出为256，表示卷积核的输入通道数为256，大小为1x1。这样就得到了mask feat和mask kernel。

然后看函数

loss_by_feat

，其中根据模型输出和ground truth计算损失。这里只看下将mask_feat和mask kernel相乘得到最终预测mask的代码如下，这里卷积核本来的个数应该是 (S^2)，但在label assignment过程中根据gt的中心点落入的grid和以及这一层level负责预测的实例大小范围得到这一层level的正样本个数为8，因此根据对应的正样本index提取出有用的8个卷积核，然后对mask_feat进行卷积得到8个实例分割的binary mask。

img_lvl_mask_pred = F.conv2d(
    img_mask_feats,  # (1,256,184,336)
    img_lvl_pos_kernel_pred.permute(1, 0).view(
        num_kernel, -1, self.dynamic_conv_size,
        self.dynamic_conv_size),  # (256,8)->(8,256,1,1)
    stride=1).view(-1, h, w)  # (1,8,184,336)->(8,184,336)
# 和v1的主要区别就在这,这里只有8个grid负责检测目标，因此只产生8个mask，而不是像v1里所有grid都生成一个mask有冗余

实验结果

在COCO测试集上和其它实例分割SOTA方法的性能对比如表1所示，可以看到SOLOv2超越了SOLOv1和其它方法。

精度-速度的对比如图1所示，可以看到SOLOv2的精度-速度trade-off是最优的。