SOLOv2 論文學習

Abstract

本文的目的是構建一個簡單、直接、快速、性能強的實例分割框架。作者延續了SOLO的方法。作者在本文中更進一步，動態地學習目標分割器的 mask head，這樣 mask head 就依賴於它的位置。此外，mask 分支被解耦成一個 mask kernel 分支和 mask 特徵分支，分別學習卷積核和卷積的特徵。

而且，作者提出了 Matrix NMS，大幅度地降低了推理時間。Matrix NMS 通過並行的 matrix 操作來執行NMS，效果要更好。本文的方法簡潔直接，在速度和精度上超過了目前SOTA方法。SOLOv2輕量版本的速度爲$31.3FPS$，AP爲$37.1\%$。此外，它也可以用於目標檢測（有了mask，順帶就有了）和全景分割，顯示了SOLOv2除了實例分割，也可以作爲許多實例級識別任務的baseline。代碼位於：git.io/AdelaiDet。

1. Introduction

一般的目標檢測方法都需要定位目標物體並識別其類別。邊框在表示目標位置方面非常簡潔，因而被廣泛採用。人們大量地研究了邊框定位目標物體，包括問題的梳理、網絡結構、後處理方式，它們都關注在邊框的優化和處理上。這些方法顯著提升了檢測性能和效率，有助於下游應用的產生。但是邊框是非常粗糙、不自然的。人類的視覺系統可以輕鬆地定位到物體邊界。實例分割通過mask來定位物體，將目標定位推到了像素級別，給實例級的感知和應用提供了可能性。截至目前，多數方法都是以邊框的視角來進行實例分割。與邊框檢測和基於邊框檢測的實例分割方法相比，單純的實例分割方法及其後處理仍然有待探索。

最近提出的SOLO有兩個像素級分類的子任務，它利用標準的FCN，大幅地簡化了實例分割的過程。SOLO 以全卷積、box-free、grouping-free 的方式，將一張圖片作爲輸入，直接輸出實例mask以及對應的類別概率。這樣我們就可以專心地研究如何輸出更好的目標masks，提升模型表現和推理速度。本文中，作者在2個方面改進了SOLO：mask 學習和mask NMS。

作者首先介紹了一個動態機制，通過位置動態地分割物體。Mask 學習可以分爲兩個部分：卷積核學習和特徵學習。網絡產生的分類器將像素點分類爲不同的位置類別。對於特徵學習，人們在語義分割中提出的方法可以用於提升模型性能。對於實例分割，受[17]中語義FPN啓發，作者構建了一個統一的、高分辨率的 mask 特徵表示。在第三章中，介紹了由SOLOv1到SOLOv2 的mask學習的推導過程。

作者進一步提出了一個高效的 matrix NMS 算法。爲了減少重複預測，在目標檢測方法的後處理階段上使用了NMS。以被普遍採用的多類別NMS爲例。對於每個類，我們根據置信度對預測邊框進行排序。然後對每個預測邊框，去除高度重疊的預測框。這項操作比較耗時。對於mask NMS，它就更加耗時了。與邊框相比，它需要更多的時間來計算每個mask對之間的IOU。作者通過matrix NMS解決了這個問題，它利用並行矩陣運算完成NMS操作。Matrix NMS 在速度和準確率上都超過了現有的NMS方法。PyTorch實現的 Matrix NMS 可以 1ms 處理500多個masks，其AP要比Fast NMS高$0.4\%$。

有了這些改進，SOLOv2 的AP要比SOLOv1高$1.9\%$，速度要快$33\%$。通過單張 V100 顯卡，在MS COCO數據集上，Res-50-FPN SOLOv2 取得了$38.8\%$的mask AP，速度爲18 FPS。SOLOv2 輕量版本的速度可以達到$31.3$ FPS，mask AP 爲$37.1\%$。有趣的是，儘管作者沒有用到邊框，作者直接將預測的mask轉化爲邊框，對於邊框目標檢測任務，它的AP也可以達到$42.4\%$，甚至超越了許多SOTA的方法。

作者認爲，簡單、快速、足夠強大的實例分割方法可以作爲目標邊框檢測方法的替代品，而SOLOv2可能會非常重要。

2. Related Works

這裏作者回顧了最近的一些有關工作。

2.1 Instance Segmentation

實例分割非常具有挑戰性。它要同時做實例級別和像素級別的預測。現有的方法可以歸納爲三大類。自上而下的方法從目標檢測的角度來解決這個問題，先檢測再分割邊框內的物體。最近的一些方法[3,38,35]基於anchor-free的檢測器FCOS進行實例分割，效果都不錯。自下而上的方法則將實例分割看作爲一個先標註再聚類的問題，學習像素級的embeddings，然後再將它們聚類成多個groups。Direct methods直接地進行實例分割，無需邊框檢測或embedding 學習。本文中，作者繼承了SOLO的設計，進一步探索了直接的實例分割方法。

作者特別地比較了YOLACT。YOLACT學習一組係數，對每個anchor box都歸一化爲$[-1,1]$。在推理時，它首先進行邊框檢測，然後用預測的邊框來裁剪組裝好的masks。BlendMask 改進了YOLACT，在速度和精度之間取得更好的平衡。

本文方法是從SOLO演進而來，直接將原始mask預測解耦成卷積核學習和特徵學習。不再需要anchor box，也不再需要歸一化，也不需要邊框檢測。作者直接將輸入圖片映射到目標類別和目標masks。訓練和推理都更加簡單。該方法取得了明顯的性能提升（$6\%$的AP提升），最佳模型取得的AP是$41.7$，而YOLACT最佳的AP是$31.2\%$。

2.2 Dynamic Convolutions

在傳統的卷積層中，學到的卷積核是固定的，獨立於輸入，對於每張圖片和圖片中的每個位置，其權重都是一樣的。之前的一些工作研究了在卷積中如何增加靈活性。Spatial Transform Networks[15]預測了一個全局參數化變換，將特徵圖變形，使得網絡根據輸入不同，自適應地變換特徵圖。Dynamic filter[16]主動地預測卷積濾波器的參數。它通過具體樣本的方式，動態地對每個圖片生成濾波器。Deformable Convolutional Networks[8] 通過預測每個圖像位置的偏移量，動態地學習採樣位置。本文中，作者在實例分割中加入了動態機制，通過位置來學習實例分割。

2.3 NMS

NMS在計算機視覺任務上被廣泛採用，成爲了目標檢測中的核心組建。最近人們提出了一些工作來改進傳統的NMS。它們可以分爲兩類，要麼提升速度，要麼改進準確率。Soft-NMS 沒有根據閾值直接去移除(hard removal)重複的預測，而是降低與較高得分預測有重疊的相鄰邊框的置信度得分。Soft-NMS 的檢測準確率要比傳統的NMS高一點，但是由於sequential operations，其推理速度要慢一些。Adaptive NMS 對每個實例使用動態的抑制閾值，專門爲人羣行人檢測而設計。爲了加速推理，[2] 中的Fast NMS 通過並行的方式來判斷預測邊框是否要保留。但是它會犧牲一定的性能。和傳統的方法相比，本文的 Matrix NMS 同時解決了hard removal和sequential operations 的問題。利用PyTorch實現的 Matrix NMS，速度可以達到1ms處理500個 masks，準確率要比 Fast NMS 高$0.4\%$。

3. A Revisit to SOLOv1

SOLOv1 的核心思想就是通過位置來分割物體。輸入圖像被分割爲$S\times S$個網格。如果物體的中心落在了某個網格內，那個網格就負責預測該物體的語義類別，並給每個像素點賦一個位置類別。它有兩個分支：類別分支和mask分支。類別分支預測語義類別，mask分支則分割物體實例。具體點，類別分支輸出$S\times S\times C$形狀的張量，其中$C$是類別個數。Mask分支產生的輸出張量爲$M\in \mathbb{R}^{H\times W\times S^2}$。$S^2$個通道中的第$k$個負責分割網格$(i,j)$的實例，其中$k=i\cdot S+j$。

我們放大看看mask分支最後一層都做了什麼。最後一層是一個$1\times 1$卷積層，輸入特徵是$F\in \mathbb{R}^{H\times W\times E}$，輸出$S^2$個通道，即張量$M$。卷積核是$G\in \mathbb{R}^{1\times 1\times E\times S^2}$。操作可以寫作：

$M=F\times G.$

這一層可以看作爲$S^2$個分類器。每個分類器負責判斷像素點是否屬於某個位置類別。

如SOLO[33]中所說，預測結果$M$有一些冗餘，因爲圖像中的大多數的物體位置都比較稀疏。在推理時，KaTeX parse error: Expected group after '^' at position 2: S^̲個分類器中，只有一小部分在實際工作。Decoupled SOLO 解決了這個問題，它將$S^2$個分類器解耦爲兩組分類器，每組$S$個，分別對應$S$個水平位置類別與$S$個垂直位置類別。因此，輸出空間就由$H\times W\times S^2$降低到了$H\times W\times 2S$。

從另一個角度看，既然輸出結果$M$是冗餘的，特徵$F$是固定的，爲什麼不直接學習卷積核$G$呢？這樣從$S^2$個分類器中，我們就可以直接選出有用的，動態地進行卷積操作。模型的參數也就減少了。由於預測卷積核是動態地取決於輸入，它就比較靈活。此外，每個分類器都取決於位置。這與SOLO的核心思想是一致的，利用位置來分割物體，而且更進一步，利用位置來預測分割器。

4. SOLOv2

4.1 Dynamic Instance Segmentation

作者繼承了SOLOv1中的一些設定，如網格單元，多層級預測，CoordConv和損失函數。作者提出了動態機制，原來的mask分支被解耦爲mask核分支和mask特徵分支，分別預測卷積核和卷積特徵。圖2爲SOLOv1的比較。

4.1.1 Mask Kernel Branch

Mask kernel 分支位於預測head內，平行的有語義類別分支。預測 head 在 FPN 輸出的特徵圖金字塔上工作。Head內的2個分支都有4個卷積層來提取特徵，和1個最終的卷積層做預測。Head 的權重在不同的特徵圖層級上共享。作者在kernel分支上增加了空間性，做法是在第一個卷積內加入了歸一化的座標，即輸入後面跟着兩個額外的通道。

對於每個網格，kernel分支預測$D$維度的輸出，表示預測卷積核的權重，其中$D$是參數的個數。爲了產生$E$輸入通道的$1\times 1$卷積的權重，$D=E$。對於$3\times 3$卷積，$D = 9E$。這些權重取決於位置，即網格單元。如果我們將輸入圖像分爲$S\times S$個網格，輸出空間就是$S\times S\times D$。對這個輸出沒有用激活函數。

4.1.2 Mask Feature Branch

Mask 特徵分支需要預測實例特徵圖$F\in \mathbb{R}^{H\times W\times E}$，其中$E$是mask特徵的維度。$F$會與mask kernel分支的輸出做卷積。如果使用所有的預測權重（即$S^2$個分類器），則輸出的實例mask將會是$H\times W\times S^2$，這與SOLOv1的輸出空間一樣。

因爲mask特徵和mask kernel 被解耦並單獨預測，我們就有2種方法來構建mask特徵分支。我們可以將它放入head內，與kernel分支平行。我們對於每個FPN層級都預測mask特徵。或者，對於所有的FPN層級預測一個統一的mask特徵表示。作者在後面的實驗中有比較這兩種實現。最終，作者使用了後者，因爲它效果更好。

爲了學習一個統一的、高分辨率的mask特徵表示，受[17]語義分割啓發，作者使用了特徵金字塔融合。在多個重複的stages之後，這些stages由$3\times 3$卷積、group norm、ReLU和$2\times$雙線性上採樣組成，FPN特徵$P2$至$P5$就被融合爲一個輸出，尺度爲$1/4$。Element-wise求和之後，最後一層包含了$1\times 1$卷積、group norm和ReLU。細節信息如圖3所示。注意在進行卷積和雙線性上採樣之前，作者將歸一化的像素座標輸入進最深的FPN層級（$1/32$大小）。這些精確的位置信息對於位置敏感性和實例特徵預測非常關鍵。

4.1.3 Forming Instance Mask

對於每個網格$(i,j)$，首先得到其mask kernel $G_{i,j,:}\in \mathbb{R}^D$。然後將$G_{i,j,:}$與$F$做卷積操作，得到實例mask。對於每個預測層級，總共會有最多$S^2$個masks。最後，作者使用Matrix NMS 得到最終的實例分割結果。

4.1.4 Learning and Inference

Label assignment和損失函數與SOLOv1是一樣的。訓練損失函數如下：

$L=L_{cate} + \lambda L_{mask}$

其中$L_{cate}$是語義類別分類的Focal Loss，$L_{mask}$是mask預測的Dice Loss。

在推理時，作者首先將圖像輸入進主幹網絡和FPN，得到網格$(i,j)$處的類別得分$\text{p}_{i,j}$。首先使用置信度閾值0.1來過濾掉低置信度的預測。然後用相應的預測mask kernels對mask特徵做卷積運算。執行完Sigmoid函數後，使用閾值0.5將預測的soft masks轉化爲二元masks。最後一步就是Matrix NMS。

4.2 Matrix NMS

動機：Matrix NMS 受到Soft NMS啓發。Soft NMS 會根據重疊度來降低其它檢測的得分，將其作爲重疊度的單調遞減函數$f(iou)$。根據IoU的值，遞歸地降低得分，可用一個極低的分數閾值來去除較高的IoU檢測。但是，這個流程是串行的，與Greedy NMS類似，無法並行實現。

Matrix NMS從另一個角度來看待這個過程，它考慮的是預測mask $m_j$ 是如何被抑制的。對於$m_j$，它的 decay factor 受如下因素影響：

• (a) 每個預測$m_i$對$m_j$的penalty ($s_i > s_j$)；
• (b) $m_i$被抑制的概率。

對於(a)，$m_i$對$m_j$的penalty可以通過 $f(iou_{i,j})$計算出來。對於(b)，$m_i$被抑制的概率計算起來就有些複雜了。但是，這個概率通常與IOU正相關。因此，作者直接用與$m_i$重疊度最高的預測來近似這個概率：

$f(iou_{.,i}) = min_{\forall s_k > s_i} f(iou_{k,i})$

最終的 decay factor 就是：

$decay_j = min_{\forall s_i>s_j} \frac{f(iou_{i,j})}{f(iou)_{.,i}}$

然後更新得分爲：$s_j = s_j \cdot decay_j$。

作者採用了2個最簡單的遞減函數，一個是線性的：

$f(iou_{i,j}) = 1-iou_{i,j}$

另一個是高斯：

$f(iou_{i,j}) = \exp(- \frac{iou^2_{i,j}}{\sigma})$

實現：Matrix NMS 所有的操作都可以單階段地實現，不需要遞歸。作者首先針對 top N 預測，計算一個$N\times N$ IOU 矩陣。對於二元mask，IOU 矩陣可以高效率地通過矩陣運算實現。然後在IOU矩陣的列上，我們就可以得到重疊度最高的IOUs。接下來，所有分數較高預測的 decay factors 就計算出來了，我們通過 column-wise min 操作來選擇每個預測的 decay factor （等式4）。最後用 decay factor 來更新分數。在使用時，我們只需要進行閾值篩選，選擇top-k得分的masks作爲最終的預測。

圖4爲Matrix NMS 的僞代碼。Matrix NMS 的速度要比傳統的NMS快9倍，而且更加準確（表7）。作者證明了Matrix NMS 可以作爲傳統NMS的替代項，在速度和精度上都很優秀，而且可以很容易地融入目前SOTA的分割/檢測方法。