IoU-aware Single-stage Object Detector for Accurate Localization 論文筆記

前言

一般來說，multi-stage目標檢測器的AP較好，而檢測效率較低，而single-stage目標檢測器雖然檢測效率較高，但AP不如multi-stage檢測器。本文就是希望在保持檢測效率的同時，提高single-stage檢測器的AP，尤其是定位精度。論文證明了single-stage檢測器中分類分數和定位精度之間的低相關性，導致了模型的性能尤其是定位精度較低。而過低的相關性主要是因爲分類子網和定位子網是相互獨立的，它們並不瞭解對方的信息。在提取anchor之後，分類子網會在不瞭解IoU（anchor與gt box）的情況下爲每個anchor預測分類分數，造成分類與定位間的mismatch問題，比如說，一些檢測結果的分類分數很高但IoU較低，而另一些檢測結果則相反。在推理時，缺失IoU會帶來兩個嚴重的問題：

• 在NMS中，所有的檢測結果基於分類分數進行排序，如果有一些檢測結果的分類分數很高但IoU較低，它肯定排在前面，從而導致那些高IoU但分類分數較低的檢測結果被過濾；
• 在計算AP時，所有的檢測結果依然按照分類分數進行排序，如果高分類分數低IoU的檢測結果排在低分類分數高IoU的檢測結果前面，那麼高IoU閾值的AP會降低。

爲了解決這個問題，本文提出了基於RetinaNet的IoU-aware single-stage目標檢測器，在迴歸分支中平行於regression head添加了一個預測IoU的head，來爲每個選中的anchor預測IoU。

• 在訓練時，IoU head與另外兩個分支一起訓練；
• 在inference時，通過將分類分數與預測的IoU相乘，得到最終的檢測分數，並用於最後的NMS和AP計算。

由於最終的檢測分數與IoU之間更具有相關性，因此前面提出的問題得以被解決，同時模型的定位精度也得到了提升。

本文的貢獻如下：

• 提出IoU-aware single-stage目標檢測器來解決single-stage檢測器中分類分數與定位精度之間的mismatch問題，在不犧牲檢測效率的情況下，模型的定位精度得以提升；
• 分析了IoU-aware single-stage檢測器爲什麼有用，性能的上限，以及預測的IoU和gt IoU之間的差別。

---

模型設計

IoU-aware single-stage檢測器對RetinaNet進行修改，backbone和FPN都是一樣的，只是在迴歸分支中添加了一個IoU prediction head，這個head和regression head的結構是平行的，分類分支依然保持不變，如上圖所示。

爲了保持模型的效率，IoU prediction head僅包含一個$3 \times 3$的卷積層，後跟一個sigmoid激活層，用於將預測的IoU控制在[0，1]之間。本文也考慮過其它方法來設計這個IoU head，比如設計一個平行於分類分支和迴歸分支的IoU預測分支，但這樣做會降低模型的效率。本文的設計方法帶來的計算開銷基本可忽略不計，並能提升模型的性能。

---

Training & Inference

training

與RetinaNet相同，採用focal loss作爲分類損失，L1 loss作爲迴歸損失，如下式（1）（2）所示。而對於IoU prediction損失，採用的是二值交叉熵損失，並且只計算正樣本的損失，如下式（3）所示。其中$IoU_i$表示爲每個detected box預測的IoU，$\hat{IoU_i}$表示正樣本與對應的gt box之間的IoU。在訓練時，IoU prediction head和其它兩個head一起共同參與訓練。

整個模型的損失爲：

inference

在inference時，通過將分類分數$p_i$與預測的$IoU_i$相乘，以得到最終的檢測分數$S_{det}$，參數$\alpha$的範圍是[0，1]，用於控制$p_i$和$IoU_i$對最終的$S_{det}$所做的貢獻。最終的檢測分數能夠很好地關聯分類分數和定位精度，比單獨使用分類分數要好。然後根據$S_{det}$對檢測結果進行排序，以進行NMS和AP計算。在排序過程中，定位錯誤而分類分數較高的檢測結果將會減少，而定位準確但分類分數較低的檢測結果將會增加，並因此提升了模型的定位精度。

---

一些討論

1. IoU-aware RetinaNet的上界

爲了評估IoU-aware RetinaNet的上界，本文將預測的IoU替換爲gt IoU，並用這個gt IoU來計算最終的檢測分數。定義檢測結果與其最近的gt box之間的IoU爲gt IoU，具體來說就是，計算檢測結果與所有類別的所有gt box之間的IoU，最大值作爲gt IoU，記爲$IoU_{truth}$。

如上表所示，當使用預測的IoU時，IoU-aware RetinaNet與RetinaNet相比能提升1.7%~1.9%的AP；而當使用gt IoU時，AP能提升14.2%~14.5%。也就是說，對於IoU-aware RetinaNet來說AP還有12.3%~12.8%的提升空間。由此可得出以下結論：

1. 很多目標已經被box檢測出來了，但是在inference時這些box由於過低的檢測分數被丟棄了；
2. 儘管預測的IoU能夠緩解檢測分數與定位精度之間的mismatch問題，但是與gt IoU相比，預測的IoU依然有較大的不足。

如果預測的IoU的準確度能夠進一步提升，那麼模型的性能也能得到很大的改善。

2. 爲什麼IoU-aware RetinaNet能夠提升模型的性能？

首先定義$IoU_{eval}$：在評估檢測結果時，計算檢測結果與屬於同一類的所有gt box之間的IoU，選擇最大值作爲$IoU_{eval}$，它用來評估檢測結果是true positive還是false positive。

分別從以下3中模型中選取10000個檢測結果，橫軸是定位精度（$IoU_{eval}$），縱軸是檢測分數（detection confidence）：

• 圖（a）：RetinaNet，檢測分數就是分類分數。可以看到，大量的檢測結果具有較高的定位精度，但檢測分數較低；
• 圖（b）：使用預測的IoU的IoU-aware RetinaNet，檢測分數是分類分數與預測的IoU相乘。可以看到，與RetinaNet相比，檢測分數略有提高。這也是爲什麼IoU-aware RetinaNet能提升模型的性能；
• 圖（c）：使用gt IoU的IoU-aware RetinaNet，檢測分數是分類分數與gt IoU相乘。可以看到，與圖（b）相比，進一步提升了檢測分數與定位精度。也就是說，預測的IoU與gt IoU之間依然有非常大的差距，因此預測的IoU以及模型的性能依然有非常大的進步空間。

---

結論

本文證明了分類分數與定位精度之間的低關聯性，是single-stage模型定位精度較低的主要原因，並由此提出IoU-aware single- stage目標檢測器來解決這個問題，具體做法是在RetinaNet的基礎上，在迴歸分支中添加了一個IoU prediction head，來預測detected box的IoU，將分類分數與預測的IoU的乘積作爲最終的檢測分數，並根據最終的檢測分數進行NMS和AP計算。實驗證明這種方法是有效的，可以提升模型的性能。除此之外，本文還證明預測的IoU與gt IoU之間依然有巨大的差距，並限制了模型性能的進一步提升。特徵對齊和attention機制可能對於IoU預測的準確度來說是非常重要的，以後的工作會提到這點。