港大&字節跳動 | 移除NMS后處理，提出端到端單階段目標檢測器（附源代碼）

目標檢測是分類和定位的聯合任務，只考慮位置信息的標簽分配和網絡優化，會導致冗余的高分辨檢測框。來自香港大學和字節跳動的研究者提出的 OneNet，首次實現了在 dense detector 中無需 NMS 后處理。

目標檢測是計算機視覺領域的基礎性任務之一，并且賦能大量的下游應用。當前目標檢測器存在的一大挑戰是標簽分配問題。特別地，如何定義每個目標的正樣本和背景的負樣本始終是一個懸而未決的難題。數十年來，目標檢測中的正樣本一直是候選框，它與真值框的 IoU 大于閾值。現代檢測器在圖像網格上預定義數千個錨框，并在這些候選框上執行分類和回歸任務。這種基于框的標簽分配方法被稱為「框分配」。

盡管多年來候選框方法在目標檢測領域占據主導地位，但檢測性能容易受到錨框大小、長寬比和數量的影響。為了消除候選框方法的手動設計和復雜計算，研究人員提出了無錨（anchor-free）檢測器。無錨檢測器中的標簽分配被簡化成了從網格點到目標框中心的點距。這種基于點的標簽分配范式被稱為「點分配」。無論是框分配還是點分配方法，都面臨著多對一（many-to-one）分配的常見困境。對于一個真值框而言，它有不止一個正樣本。檢測性能容易受到分配過程中超參數的影響。更糟糕的是產生了冗余和近似重復的結果，導致非最大抑制（NMS）成為了必要的后處理。

最近，一對一（one-to-one）分配在稀疏候選和多階段細化檢測器中取得了成功，其中一個真值邊界框只分配給一個正樣本，其余皆為負樣本。這類方法不需要 NMS 即可直接輸出檢測結果。檢測器的檢測準確率性能也非常好。但是，密集候選和單階段細化檢測器可能會更快和更簡單。所以為何不為端到端單階段檢測器設計一種直接的一對一標簽分類策略呢？

在本文中，來自香港大學和字節跳動 AI 實驗室的研究者發現：標簽分配中樣本與真值之間缺乏分類代價是 one-stage 檢測器移除非最大抑制（NMS）并實現端到端的主要障礙。現有的 one-stage 目標檢測僅通過位置代價來分配標簽，例如框 IoU、點距離。在缺少分類代價的情況，單獨的位置代價將將導致高置信度得分在推理中產生冗余框，從而使 NMS 成為必要的后處理。

具體而言，為了設計一種端到端 one-stage 目標檢測器，研究者提出了最小代價分配（Minimum Cost Assignment）。代價是樣本與真值之間的分類代價和位置代價的總和。對于每個目標真值，僅將一個最小代價樣本分配為正樣本，其他都是負樣本。為了評估該方法的有效性，研究者設計了一個非常簡單的 one-stage 檢測器 OneNet。

實驗結果表明，在經過「最小代價分配」訓練后，OneNet 避免了產生重復框，并實現了端到端目標檢測。在 COCO 數據集上，OneNet 實現了 35.0 AP/80 FPS 和 37.7 AP/50 FPS，圖像大小為 512 像素。

最小代價分配

最小代價分配旨在解決端到端 one-stage 目標檢測的標簽分配，此前 one-stage 檢測器的標簽分配方法有框分配和點分配。該研究提出了全新的最小代價分配。

最小代價分配是一種直接的方法，對于每個真值，在所有樣本中僅選擇一個最小代價樣本作為正樣本，其余都是負樣本。該方法不涉及手動制定的啟發式規則或者復雜的二分圖匹配。以下算法 1 給出了最小代價分配的說明性示例：交叉熵損失作為分類代價，L1 損失作為位置代價。

在密集探測器中，分類損失為焦點損失。繼 [2, 42, 31] 之后，位置代價包含 L1 損失和泛化 IoU（GIoU）損失。最后，代價如下等式所示：

其中 C_cls 是預測分類和真值類標簽的焦點損失，C_L1 和 C_giou 分別是歸一化中心坐標與預測框和真值框高度和寬度之間的 L1 損失和 GIoU 損失。λ_cls、λ_L1 和λ_giou 是每個分量的系數。

OneNet 目標檢測器

OneNet 是一種基礎的全卷積 one-stage 檢測器，沒有任何后處理（如 NMS）。其流程圖如下所示：

圖 3：OneNet 的流程。

主干

OneNet 的主干是一個先自下而上，然后自上而下的結構。自下而上的部分是 ResNet 架構[13]，用于生成多尺度特征圖。具有橫向連接（FPN）的自上而下架構的作用是生成用于目標識別的最終特征圖。輸出特征的形態是 H/4xW/4xC，其中 H 和 W 分別是輸入圖像的高度和寬度。

頭（head）

頭通過兩個并行卷積層在特征圖 H/4xW/4 的每個網格點上進行分類和定位。分類層預測 K 個目標類別在每個網格點上存在目標的概率。位置層可以預測從每個網格點到真值框四個邊界的偏移量。

訓練

標簽分配是最小代價分配。訓練損失類似于匹配代價，包括焦點損失、L1 損失和 GIoU 損失。

推理

最終輸出是直接的前 k 個（例如 100 個）得分框，沒有任何后處理流程（如 NMS 或最大池操作）。

多頭訓練策略

在實現過程中，研究者采用了一種可選擇的多頭訓練策略，主要包括級聯預測頭和權重共享機制。其示意圖如下所示：

多頭訓練策略。

級聯預測頭在訓練過程中使用，其中分類卷積和回歸卷積的參數分別在頭中共享。推理過程中僅使用 first head，因此推理速度不會產生延遲。

實驗

正樣本的可視化

該研究將 CenterNet [41] 和 OneNet 正樣本進行了可視化，如下圖 5 所示。兩種方法最大的區別在于 CenterNet 遵循標簽分配位置代價（ location cost），而 OneNet 在分類代價和位置代價中遵循最小代價分配。

圖 5：正樣本的可視化。第 1 行是位置代價。第 2 行是分類代價與位置代價。正網格點由圓圈突出顯示，畫出來的邊框為真值框。僅有位置代價分配的正樣本是最接近真值框中心的網格點。添加分類代價，使得正樣本成為更具識別區域的網格點。例如圖 5 中斑馬的頭部。

控制變量研究

標簽分配的研究。在下表 1 中，研究者進行了一系列的實驗來研究標簽分配對結果的影響。預定義位置代價是特征圖中柵格點的固定位置與地面真值框中心位置之間的距離。被預測框的位置代價和分類損失參見第 3 節。分類代價是移除 NMS 的關鍵。

下表 2 為多頭訓練的控制變量實驗。結果表明多頭訓練可以使用較大的學習率，得到較高的準確率。同時，多頭訓練與單頭推理獲得了較高的準確率，以及和基準相似的推理速度。

下表 3 研究了圖像尺寸對結果的影響。由結果可知，較大的圖像尺寸具有較高的準確率，但是，推理速度卻慢了。

下表 4 為訓練 epoch 對結果的影響。檢測準確率的增長速度在訓練 epoch 為 144 時開始下降。因此，該研究選擇 epoch 為 144 作為基準配置。

THE END