SOLOv2算法学习

SOLOv2神经网络学习

SOLOv2是在SOLO的基础上进行的改进，原文为《SOLOv2: Dynamic and Fast Instance Segmentation》。代码在https://git.io/AdelaiDet上可以找到。

摘要

SOLOv2相对SOLO的主要改动有两点：

（1）新框架由一个实例掩码表示方案来赋能，能够动态分割图像中的每个实例，不用借助bbox检测。具体来说，将目标掩码的生成（Mask branch）解耦成掩码核预测（Mask kernel branch）和掩码特征学习（Mask feature branch），分别负责生成卷积核和需要卷积的特征映射。

（2）SOLOv2通过作者新的矩阵非最大值抑制（NMS）技术显著减少了推理运算的开销。

Introduction

实例分割任务可以分解为像素级分类的两个子任务，SOLO算法以图像为输入，直接输出实例掩码和相应的类概率，采用完全卷积、无框、无组的范式。然而，限制SOLO性能的主要瓶颈有三个：低效的掩码表示和学习；分辨率不够高，无法进行更精细的掩膜预测；mask NMS较慢。而在SOLOv2中，三个问题都得到了解决。

首先，作者引入了一种动态的方案，可以根据位置动态分割对象。具体来说，掩码学习过程可以分为两个部分：卷积核学习和特征学习。在将像素分类到不同位置类别时，网络会动态预测掩码内核，并以输入为条件进行适应调整。进一步地，作者构建了一个统一的、高分辨率的掩码特征表示方法，用于实例感知的分割。

作者还提出了一种高效的矩阵NMS算法。作为已知重复预测的后处理步骤，非最大值抑制（NMS）作为最先进的对象检测系统的一个组成部分。

以广泛采用的多类别NMS为例，对于每个类别，预测按可信度降序排序，对于每个预测，删除所有其他高度重叠的预测。顺序和递归操作导致不可忽略的延迟，对于mask NMS，这个缺点被进一步放大。与bbox对比，计算每个mask pair的IoU需要更多的时间。

作者通过引入矩阵NMS解决这个问题，使用并行的矩阵操作一次性执行NMS。矩阵NMS能在不到1ms的时间内处理500个mask，并且比最近提出的快速NMS的性能提高0.4%的AP。

Related Work

实例分割

实例分割的挑战性在于其同时进行实例级和像素级的预测，现有的方法主要有两种。

自上而下的方法从对象检测的角度来解决问题，即先检测、再分割框内的对象。

自底向上的方法将任务视为一个先对每个像素进行标签分类，再进行集群的问题，比如学习逐像素嵌入，然后把它们聚类成组。

作者将SOLOv2与YOLACT进行了比较，YOLACT学习一组系数，对于每个锚盒，将这些系数归一化到$(-1,1)$。在推理过程中，它首先进行边界框检测，然后使用预测的框来裁剪组装好的掩膜。而SOLOv2通过直接将原始的掩膜预测解耦到核学习和特征学习，不需要锚盒、归一化和bbox。

---

SOLOv2详解

实例分割系统应该在像素级别分离不同的实例。为了区分实例，SOLO将输入图像在概念上分为$S\times S$的网格。如果一个对象的中心落入一个网格单元，那么该网格单元对应于该对象的二进制掩码。因此，系统共输出$S^2$个masks，记为$M\in R^{H\times W\times S^2}$。第$k$个通道负责对位置$(i,j)$的实例进行分割，$k=i\cdot S+j$。

然而，SOLO有三个主要的瓶颈：

（1）低效的掩码表示方法和学习。预测具有$S^2$通道的输出张量$M$需要大量的内存和计算。另外，由于不同的FPN层的S是不同的，所以每个级别的最后一层是单独学习的，而不是共享的，导致了训练低效。

（2）不准确的mask预测。更好的预测需要高分辨率的遮罩来处理物体边界的细节，但大分辨率将大大增加计算成本。

（3）mask NMS较慢。与box NMS相比，mask NMS需要更多的时间与更大的开销。

动态实例分割

首先，回顾下SOLO中mask的生成，为了生成$S\times S$网格对应的$S^2$通道的实例掩码，最后一层以一层金字塔特征$F\in R^{H\times W\times E}$作为输入，最后应用到一个具有$S^2$输出通道的卷积层。该操作可以写成：
$$M_{i,j}=G_{i,j}\ast F$$
其中，$G_{i,j}\in R^{1\times 1\times E}$是卷积核，$M_{i,j}\in R^{H\times W}$是只包含一个中心位于位置$(i,j)$的实例的最终掩码。

换句话说，需要两个输入$F$和$G$来生成最终的掩码$M$。之前的工作明确地输出整个$M$用于训练和推理，注意张量$M$非常大，直接预测$M$内存和计算效率低下。在大多情况下，物体在图像中的位置是稀疏的，$M$是冗余的，因为$S^2$内核中只有一小部分在一个推断过程中起作用。

从另一个角度来说，单独学习$F$和$G$，可以简单从预测的$S^2$内核中选取有效的核，并动态地进行卷积，模型参数量会减少。而且，由于预测的内核是根据输入动态生成的，因此它受益于灵活性和自适应特性。另外，每个$S^2$内核的位置都是有条件的，它符合按位置进行对象分割的核心思想，并进一步地通过位置来预测分割对象。

Mask kernel G

在已知主干和FPN的前提下，在每个金字塔层预测掩码内核$G$。我们首先将输入特征$F_I\in R^{H_I\times W_I\times C}$调整为$S\times S\times C$的形状，然后使用$4\times convs$和最终的$3\times 3\times D$的卷积生成核$G$。

对于每个网络，内核分支预测$D$维输出来表示预测的卷积核权值，其中$D$是参数的个数。对于一个有$E$个输入通道的$1\times 1$卷积，$D=E$；对于$3\times 3$的卷积， $D=9E$。这些生成的权重取决于位置，也就是网格单元。如果我们将输入图像划分为$S\times S$网格，输出空间为$S\times S\times D$，输出上无激活函数。

Mask Feature F

由于mask feature和mask kernel分别预测，因此有两种构造mask feature的方法，我们可以把它和内核分支一起放到head中。这意味着我们预测每个FPN层的mask feature。或者，预测所有FPN层的统一mask feature表示方法。最终采用了后者。

形成实例mask

对于位于$(i,j)$的每个网格单元，首先得到mask kernel $G_{i,j}\in R^D$。然后，$G_{i,j}\in R^D$与$F$进行卷积，总的来说，每个预测级别最多有$S^2$个mask。最终，利用矩阵NMS得到最终的实例分割结果。

学习与推理

训练损失函数定义如下：$$L=L_{cate}+\lambda L_{mask}$$

其中，$L_{cate}$是语义分类的常规Focal Loss，$L_{mask}$是遮罩预测的Dice Loss。

在推理过程中，我们通过骨干网络和FPN转发输入图像，得到单元格$(i,j)$的类别得分$p_{i,j}$。一开始，使用0.1的置信度阈值来过滤低置信的预测。然后使用相应的预测掩码核对mask特征进行卷积。在sigmod操作后，使用0.5的阈值将预测的软掩码转换为二进制掩码。最后一步则是矩阵NMS。