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是最优的。

03-25
### SOLOv2 模型简介 SOLOv2 是一种基于位置的实例分割算法,其核心思想是通过预测每个像素的位置来完成对象分割任务。相比于传统的两阶段方法(如 Mask R-CNN),SOLOv2 提供了一种更高效的单阶段解决方案[^3]。 以下是使用 SOLOv2 进行模型训练、推理以及可视化的具体实现方式: --- ### 安装依赖环境 为了运行 SOLOv2代码,需先安装必要的 Python 库和框架。通常情况下,推荐使用 Anaconda 创建虚拟环境,并按照官方文档的要求安装 PyTorch 和 MMDetection 工具包。 ```bash pip install mmcv-full==latest+torch1.x.cuda11.3 -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.x/index.html git clone https://github.com/open-mmlab/mmdetection.git cd mmdetection pip install -e . ``` 上述命令用于克隆 MMDetection 并将其作为开发模式安装到环境中[^4]。 --- ### 配置文件准备 MMDetection 中提供了多种预定义配置文件,其中包含了不同数据集和网络结构的具体参数设置。对于 SOLOv2 而言,默认配置位于 `configs/solov2` 文件夹下。 假设我们使用的配置文件名为 `solov2_light_r50_fpn_3x_coco.py`,则可以通过以下路径找到它: ```plaintext ./configs/solov2/solov2_light_r50_fpn_3x_coco.py ``` 如果需要自定义超参或者调整 backbone 结构,则可以直接修改该配置文件的内容。 --- ### 训练过程 启动分布式训练时,可利用多张 GPU 协同工作以加速收敛速度。下面是一个典型的训练脚本示例: ```bash CUDA_VISIBLE_DEVICES=0,1,2,3 PORT=29500 ./tools/dist_train.sh configs/solov2/solov2_light_r50_fpn_3x_coco.py 4 ``` 此命令表示启用四块显卡执行训练操作,同时指定端口号为 29500 来协调进程间通信[^1]。 --- ### 推理可视化 当训练完成后,保存下来的权重文件可用于后续测试环节。此时可通过如下指令加载已有的 checkpoint 对新图片做推断处理: ```bash CUDA_VISIBLE_DEVICES=0,1 bash tools/dist_test.sh ${CONFIG_FILE} ${MODEL_PATH} ${GPU_NUM} --eval bbox segm --show-dir ./test_result --show ``` 在此过程中,`${CONFIG_FILE}` 表示所选配置文件路径;`${MODEL_PATH}` 则指向最终生成的 `.pth` 格式的权值档案名;而 `${GPU_NUM}` 明确指定了参运算的图形处理器数量。 另外需要注意的是,在实际部署场景里,非极大抑制 (Non-Maximum Suppression, NMS) 可能成为性能瓶颈之一。因此优化这部分逻辑显得尤为重要[^2]。 --- ### 实现细节解析 #### Grid-to-Mask Mapping SOLOv2 将输入图像划分为固定大小的小区域(grid cells)。每一个 cell 不仅负责检测目标是否存在及其类别标签,还会生成相应尺寸的二进制掩码图谱。这种设计巧妙地避开了传统 bounding box 边界框定位难题,从而实现了更加精准的对象轮廓描绘功能。 下面是简化版的核心代码片段展示如何构建此类映射关系: ```python import torch.nn as nn class Solov2Head(nn.Module): def __init__(self, num_classes, in_channels): super(Solov2Head, self).__init__() # Define layers here... def forward(self, feats): cate_preds = [] kernel_preds = [] for idx, feat in enumerate(feats): ins_pred_x = self.ins_convs(feat) kernel_pred = self.kernel_convs(ins_pred_x) cate_pred_y = self.cate_convs(feat) cate_pred = self.solo_cate(cate_pred_y) kernel_preds.append(kernel_pred) cate_preds.append(cate_pred) return cate_preds, kernel_preds ``` 以上模块展示了分类分支 (`cate_preds`) 和核生成分支 (`kernel_preds`) 的基本架构形式。 ---
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

00000cj

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值