之前在一次組會上，師弟訴苦說他用 UNet 處理一個病灶分割的任務，但效果極差，我看了他的數據后發現，那些病灶區域比起整張圖而言非常的小，而 UNet 采用的損失函數通常是逐像素的分類損失，如此一來，網絡只要能夠分割出大部分背景，那么 loss 的值就可以下降很多，自然無法精細地分割出那些細小的病灶。反過來想，這其實類似於正負樣本極不均衡的情況，網絡擬合了大部分負樣本后，即使正樣本擬合得較差，整體的 loss 也已經很低了。
發現這個問題后，我就在想可不可以先用 Faster RCNN 之類的先檢測出這些病灶區域的候選框，再在框內做分割，甚至，能不能直接把 Faster RCNN 和分割部分做成一個統一的模型來處理。后來發現，這不就是 Mask RCNN 做的事情咩～囧～
這篇文章，我們就從無到有來看看 Mask RCNN 是怎么設計出來的。
回顧 Faster RCNN
首先，我們簡單回憶一下 Faster RCNN 的結構，看看如何針對它進行拓展。上面這個框架圖中，虛線框內就是 Faster RCNN 的大致結構了。算法過程可以粗略分為以下幾步：
- 將圖片輸入 CNN 中，得到 feature map；
- 用一個 RPN 網絡在 feature map 提取出候選框（region proposals）。這一步對應 RPN 網絡分支；
- 用另一個 CNN 進一步對該 feature map 進行特征提取，結合候選框得到很多 RoI，然后在每個 RoI 內用 RoI Pooling 提取特征，之后接上 FC 層分別預測框內物體的類別以及做 bounding box 微調。這一步對應 Fast RCNN 網絡分支。
整個流程其實非常簡潔。既然我們是想在候選框內進一步做分割，那么很自然的想法，就是在原 Faster RCNN 選出來的 RoI 中，根據 classification score 和 bounding box regression 選出得分最高的 RoI，並對這些 RoI 的框進行微調。這樣，這些 RoI 就是最可能包含物體，同時定位也更為准確的 RoI 了，之后繼續在這些 RoI 內做分割即可。另外，這個想法還能沿用之前得到的 feature map，避免卷積重復計算，可以說是一舉兩得。
事實上，Mask RCNN 采用的也是這一套思路。這種先檢測物體再做分割的兩步走策略，江湖人稱 two stage，而 UNet 這種一步到位的分割方法，則被稱為 one stage。顯然，一步到位的 UNet 實現起來簡單，大部分情況下效果也還行，但論精度，還是 Mask RCNN 更勝一籌，畢竟有候選框作為先驗支撐。
這里要注意另一個問題，雖然 UNet 和 Mask RCNN 都是處理分割，但前者又稱為 semantic segmentation，后者稱為 instance segmentation。兩者的區別可以用下面這張圖體現，semantic segmentation 在分割的時候對同一類物體一視同仁，而 instance segmentation 則需要把每個個體都單獨分割出來。因此，把 UNet 和 Mask RCNN 進行對比其實不太公平。instance segmentation 由於需要單獨分割每個個體，因此基本上所有針對 instance segmentation 的方法都需要先用一個候選框把物體找出來，之后再分割。
下面，我們就來看看 Mask RCNN 是怎么在 Faster RCNN 的基礎上實現分割的。
Mask RCNN
Mask RCNN 的整體流程圖可以參考文章開頭那個框架圖。它在 Faster RCNN 的基礎上，延伸出了一個 Mask 分支。根據 Faster RCNN 計算出來的每個候選框的分數，篩選出一大堆更加准確的 RoI（對應圖中 selected RoI），然后用一個 RoI Align 層提取這些 RoI 的特征，計算出一個 mask，根據 RoI 和原圖的比例，將這個 mask 擴大回原圖，就可以得到一個分割的 mask 了。
RoI Align 之后打算新開一篇文章針對代碼細講，本文只稍微提及 RoI Align 背后的機理，暫且就將它當作一個黑盒。
現在，我們從零出發，假設讓我們來設計 Mask RCNN，這中間每一步要如何操作，以及會面臨哪些問題。
RoI 如何到 Mask
首先第一個問題是如何處理這些 RoI 的特征，輸出分割的 mask。在 UNet 和 FCN 中，是先將圖片通過卷積操作得到 feature map，再通過 deconvolution（或者 upsample + convolution）的方法將 feature map 的尺寸還原成原圖大小。因此，我們也可以借用同樣的思路，先根據 RoI 的尺寸換算回原圖，看看這個 RoI 對應到原圖上的尺寸有多大，再將 RoI 內的 feature map 上采樣成對應尺寸的 mask，然后接一個 FCN 網絡將 mask 的通道數處理成 \(K\) 即可（假設總共有 \(K\) 種類別）。由於上采樣是可以求導的，因此反向傳播依然有效。
寫到這里，我突然覺得，Fast RCNN 中提出的 RoI Pooling 是不是也可以換成這種上/下采樣 + 卷積的方式來得到一個固定大小的 feature map，這樣，之后的 FC 層的維度也是可以匹配上的。而且 RoI Pooling 本質上也類似一種下采樣的操作，只不過采樣的方式是取鄰域中最大的數值。但轉念一想，這其實就是在問 Pooling 操作能不能用下采樣來代替，從大佬們設計的網絡結構中普遍采用 Pooling 而不是下采樣來看，應該是 Pooling 的效果會更好。這個問題就此打住，暫且認為 Pooling 的效果好於下采樣。
可以看出，我們設計的這個 mask 分支跟 UNet 或 FCN 的思路其實一樣。不過，直接下采樣可能會導致 feature map 中一些重要的信息丟失，因此，我們可以沿用 RoI Pooling 的思路，將 RoI 內的那塊 feature map 處理成指定大小的 feature map，然后采用 conv/deconv 的方法進一步轉換成 \(K \times H \times W\) 的 mask，其中 \(K\) 表示物體的類別。最后把這個 mask 按照 RoI 對應的 bounding box 換算回原圖中即可。這一步在訓練上也可以類比 FCN，先根據 bounding box 找到 ground truth 中對應的那塊 mask，再按比例縮成跟網絡輸出的 mask 一樣大小，然后根據分類損失或者 MSE 構造 Loss 函數即可。
大部分人都能走到這一步，但也就僅僅走到這一步而已。這個流程簡單直接，而且也能 work，但實操后會發現效果不佳。這里就涉及到論文中提及的 misalignment 問題。
事實上，分割對模型精度的要求比分類以及檢測要高得多，因為前者需要逐像素的標注類別信息。這意味着，如果 RoI 中的 feature map 跟原圖中對應的區域存在偏差，就可能導致計算出來的 mask 跟 ground truth 是對不齊的。我們用一張圖來說明：
原圖中的小男孩有兩個 bounding box 框住他，我們用卷積操作對圖像進行 downsample，然后根據 feature map 和原圖的比例，推算出這兩個 bounding box 對應在 feature map 上的位置和大小。結果，很不幸這兩個框的位置四舍五入后剛好對應同一塊 feature map。接着，RoI Pooling 和 FCN 會對這塊 feature map 進行處理得到 mask，再根據 ground truth 計算 Loss。兩個框對應的 ground truth 當然是不一樣的。這個時候，網絡就左右為難，同樣一個 feature map，居然要擬合兩個不同的結果，它左右為難一臉懵逼，直接導致模型無法收斂。
這就是所謂的 misalignment。問題的根源在於 bounding box 縮放時的取整。當然，RoI Pooling 本身在 pooling 的時候也是存在取整誤差的。
既然問題出在取整上，那么，很自然想到的解決思路就是放棄取整，直接根據推算得到的浮點數來處理 bounding box。如此一來，bounding box 對應到的 feature map 上就會有一些點的坐標不是整數，於是，我們需要重新確定這些點的特征值。而這一步也是 RoI Align 的主要工作。其具體的做法是采用雙線性插值，根據相鄰 feature map 上的點來插值一個新的特征向量。如下圖所示：
圖中，我們先在 bounding box 中采樣出幾個點，然后用雙線性插值計算出這幾個點的向量，之后，再按照一般的 Pooling 操作得到一個固定大小的 feature map。具體的細節，之后開一篇新的文章介紹。
論文中通過 RoI Align 和 FCN 將 RoI 內對應的 feature map 處理成固定大小的 mask（\(K \times m \times m\)，\(K\) 表示分割的類別數目），然后將該 mask 還原回原圖后，就可以得到對應的分割掩碼了。
Loss 的設計
在損失函數的設計方面，除了原本 Faster RCNN 中的分類損失和 bounding box 回歸損失外，我們還需要針對 mask 分支設計一個分割任務的損失函數。最容易想到的函數自然是 FCN 和 UNet 中用到的 Softmax + Log 的多分類損失。如下圖所示：
Mask 中的每個點都是一個 \(K\) 維的向量，我們把 ground truth 中對應的那個 mask 也縮放到 \(m \times m\) 大小，然后就可以針對每個點的向量做多分類損失。
不過，作者在做實驗的時候估計是發現這種方式訓練的網絡收斂不好，進而發現這個損失函數會出現所謂的 class competition 的問題。
class competition，顧名思義，就是不同類別之間存在競爭關系，這種競爭關系直接導致的結果就是網絡在訓練過程中，回傳的梯度存在前后不一致的地方。
打個比方，在 Faster RCNN 做 object detection 的時候，已經把某一塊 RoI 識別為汽車，但這個 RoI 內可能存在其他物體的一部分，因此分割的 mask 中，除了要將汽車分割出來外，還要把另外那個物體也分割出來。這就導致這樣的情況，在 object detection 的分支中，這塊 RoI 整體被識別為汽車，但在 segmentation 的時候，這塊 RoI 一部分被識別為汽車，一部分又要當作其他物體，如此一來，這兩個分支回傳到前面的梯度多少存在沖突，而前面的特征提取網絡可是共享的，結果網絡在學習的時候就可能出現左右為難的情況。不然，單純從 Mask 分支來看，feature map 上每個點（包括 RoI Align 插值的點）本來就和 ground truth 是一一對應的，彼此之間又哪有 competition 之說呢？當然，以上只是我看論文時的想法，並沒有做具體的實驗，所以也不一定正確。
之后我又想，在 object detection 的時候，那些 bounding box 本身就是有重疊的，換句話說，RoI 之間也是有重疊的，如果兩個 RoI 被識別為不同的物體，那么重疊那部分不也是沖突的嗎？這個時候應該找個例子看看重疊那部分的特征圖是什么樣子的。不過，我個人的想法是，網絡對這些重疊的區域可能會起到抑制作用。比如說，一輛汽車前面被一輛自行車擋着，那么汽車的 bounding box 多少會覆蓋到自行車，而自行車的 bounding box 也多少包含了汽車的一部分，但這個交集相比各自的 bounding box 而言，可能不是主體作用，網絡在對汽車的 RoI 做識別的時候，更多的會把注意力放在非重疊的那部分汽車上，而重疊那部分，雖然有一點點汽車的東西，但由於有自行車的遮擋，起到的作用不會太大。
最后需要聲明一下，這個想法完全是我個人瞎猜的，並沒有做實驗證明~.~
既然多分類效果不好，那我們就嘗試二分類。如下圖所示：
二分類的話，我們只考慮一種類別，比如，如果 ground truth 中標記了這個 bounding box 中是個人的話，那我們就只針對人的 mask 進行分割，而對這個 bounding box 中其他可能存在的物體一律忽視。上圖中，假設人的類別是第 \(K\) 類，那么，我們就只在第 \(K\) 個 mask 上和 ground truth 中人的 mask 做 sigmoid 的二分類損失。如果某個點被標注是人的一部分，就識別為 1，否則全部識別為 0。
有同學可能會有疑惑，二分類只考慮一種物體，而把其他物體的部分直接忽略掉，如果出現一種極端情況，比如有一個人的一只手出現在一個汽車的 bounding box 中，然后網絡計算這個人的 RoI 的時候，bounding box 剛好沒有把這只手包含進來，而在汽車那個 RoI 里面又不會對這只手做分割，這樣的話，這只手不就直接漏掉了嗎？確實，這種情況下，這只手會被直接忽略。不過，這種情況屬於 bounding box regression 沒有做好，因此不在本文討論范圍內。
總結一下，損失函數可以表示為：
\(L_{cls}\) 和 \(L_{box}\) 是 Faster RCNN 中的損失函數，而 \(L_{mask}\) 則是 mask 分支中的 sigmoid 二分類損失。
特征提取
特征提取部分其實可以有多種選擇，具體哪種選擇好，可能要依據具體的任務來確定。論文嘗試了 ResNet、FRN、ResNeXt 等網絡。這一部分我沒有去細究，因為這里變數比較大，針對不同的場景可以適當調整，因此這一塊就不細談了。
訓練和預測
訓練階段依賴 Faster RCNN 的輸出結果。首先根據 Faster RCNN 找出一大堆 RoI，再根據 classification score 對這些 RoI 進行排序，選出分類分數最高的前 \(N\) 個 RoI，然后根據 ground truth 中的 mask 和這個 RoI 取個交集，這個交集作為 \(L_{mask}\) 的 target。實際預測的時候，同樣先根據 Faster RCNN 選出前 100 個分數最高的 RoI，然后計算每個 RoI 的 mask。不過，由於這些 mask 是根據二分類損失訓練出來的，因此，我們要根據 Faster RCNN 提供的每個 RoI 的類別，在 mask 中找出對應的那一層作為最終分割的結果。
從這個訓練過程也可以找出一些不足的地方。比如，挑選 RoI 送入 mask 分支那一步，這個挑選的結果完全依賴於 Faster RCNN 計算的分數，一旦 Faster RCNN 出了差錯，給一些很重要的 RoI 打了很低的分，那么這些 RoI 就可能被忽略掉，之后分割就沒它們什么事了。因此，有人提出了一些改進，認為應該對這個篩選的打分機制進行修改，不應該完全依賴於 Faster RCNN 的結果，比如，Mask Scoring RCNN 就在打分中加入了 ground truth 的 mask 的 IoU 分數，從而把那些容易被忽略的 RoI 找出來。這有點像是難樣本挖掘了。
實驗
何凱明在論文中一直強調 Mask RCNN 是 without bells and whistles，意思就是 Mask RCNN 的算法中沒有什么花里胡哨的東西，都是實打實的干貨，無需特殊的調參技巧，經得起時間的考驗。為此論文中還提供了很多對比實驗來一一驗證每個模塊的作用。
-
首先是 RoIAlign 和 RoIPooling 的對比：
在 instance segmentation 和 object detection 上都有不小的提升。如此看來，RoIAlign 其實就是一個更加精准的 RoIPooling，把前者放到 Faster RCNN 中，對結果的提升應該也會有幫助。
-
sigmoid 和 softmax 的對比
這里同樣可以取得不小的提升。
-
特征網絡的選擇
總體來說，加上 FPN 網絡的效果會更好，可能因為 FPN 綜合考慮了不同尺寸的 feature map 的信息，因此能夠把握一些更精細的細節。
-
另外，論文還針對人體關鍵點檢測做了一個實驗，來體現 Mask RCNN 框架的通用性，這部分內容我還不太熟悉，就先略過了。
總結
總的來說，Mask RCNN 這種先檢測物體，再分割的思路，簡單直接，在建模上也更有利於網絡的學習。而其中，我認為兩個最重要的改進點，分別是 RoIAlign 和采用 sigmoid 二分類損失。這兩個改進的目標都是讓網絡在學習的時候能保持一致性，使得輸入到輸出之間的映射關系更加簡單直接。