港大&字節(jié)跳動(dòng) | 移除NMS后處理，提出端到端單階段目標(biāo)檢測器（附源代碼）

目標(biāo)檢測是分類和定位的聯(lián)合任務(wù)，只考慮位置信息的標(biāo)簽分配和網(wǎng)絡(luò)優(yōu)化，會(huì)導(dǎo)致冗余的高分辨檢測框。來自香港大學(xué)和字節(jié)跳動(dòng)的研究者提出的 OneNet，首次實(shí)現(xiàn)了在 dense detector 中無需 NMS 后處理。

目標(biāo)檢測是計(jì)算機(jī)視覺領(lǐng)域的基礎(chǔ)性任務(wù)之一，并且賦能大量的下游應(yīng)用。當(dāng)前目標(biāo)檢測器存在的一大挑戰(zhàn)是標(biāo)簽分配問題。特別地，如何定義每個(gè)目標(biāo)的正樣本和背景的負(fù)樣本始終是一個(gè)懸而未決的難題。數(shù)十年來，目標(biāo)檢測中的正樣本一直是候選框，它與真值框的 IoU 大于閾值?，F(xiàn)代檢測器在圖像網(wǎng)格上預(yù)定義數(shù)千個(gè)錨框，并在這些候選框上執(zhí)行分類和回歸任務(wù)。這種基于框的標(biāo)簽分配方法被稱為「框分配」。

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

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

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

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

實(shí)驗(yàn)結(jié)果表明，在經(jīng)過「最小代價(jià)分配」訓(xùn)練后，OneNet 避免了產(chǎn)生重復(fù)框，并實(shí)現(xiàn)了端到端目標(biāo)檢測。在 COCO 數(shù)據(jù)集上，OneNet 實(shí)現(xiàn)了 35.0 AP/80 FPS 和 37.7 AP/50 FPS，圖像大小為 512 像素。

最小代價(jià)分配

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

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

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

其中 C_cls 是預(yù)測分類和真值類標(biāo)簽的焦點(diǎn)損失，C_L1 和 C_giou 分別是歸一化中心坐標(biāo)與預(yù)測框和真值框高度和寬度之間的 L1 損失和 GIoU 損失。λ_cls、λ_L1 和λ_giou 是每個(gè)分量的系數(shù)。

OneNet 目標(biāo)檢測器

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

圖 3：OneNet 的流程。

主干

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

頭（head）

頭通過兩個(gè)并行卷積層在特征圖 H/4xW/4 的每個(gè)網(wǎng)格點(diǎn)上進(jìn)行分類和定位。分類層預(yù)測 K 個(gè)目標(biāo)類別在每個(gè)網(wǎng)格點(diǎn)上存在目標(biāo)的概率。位置層可以預(yù)測從每個(gè)網(wǎng)格點(diǎn)到真值框四個(gè)邊界的偏移量。

訓(xùn)練

標(biāo)簽分配是最小代價(jià)分配。訓(xùn)練損失類似于匹配代價(jià)，包括焦點(diǎn)損失、L1 損失和 GIoU 損失。

推理

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

多頭訓(xùn)練策略

在實(shí)現(xiàn)過程中，研究者采用了一種可選擇的多頭訓(xùn)練策略，主要包括級(jí)聯(lián)預(yù)測頭和權(quán)重共享機(jī)制。其示意圖如下所示：

多頭訓(xùn)練策略。

級(jí)聯(lián)預(yù)測頭在訓(xùn)練過程中使用，其中分類卷積和回歸卷積的參數(shù)分別在頭中共享。推理過程中僅使用 first head，因此推理速度不會(huì)產(chǎn)生延遲。

實(shí)驗(yàn)

正樣本的可視化

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

圖 5：正樣本的可視化。第 1 行是位置代價(jià)。第 2 行是分類代價(jià)與位置代價(jià)。正網(wǎng)格點(diǎn)由圓圈突出顯示，畫出來的邊框?yàn)檎嬷悼?。僅有位置代價(jià)分配的正樣本是最接近真值框中心的網(wǎng)格點(diǎn)。添加分類代價(jià)，使得正樣本成為更具識(shí)別區(qū)域的網(wǎng)格點(diǎn)。例如圖 5 中斑馬的頭部。

控制變量研究

標(biāo)簽分配的研究。在下表 1 中，研究者進(jìn)行了一系列的實(shí)驗(yàn)來研究標(biāo)簽分配對(duì)結(jié)果的影響。預(yù)定義位置代價(jià)是特征圖中柵格點(diǎn)的固定位置與地面真值框中心位置之間的距離。被預(yù)測框的位置代價(jià)和分類損失參見第 3 節(jié)。分類代價(jià)是移除 NMS 的關(guān)鍵。

下表 2 為多頭訓(xùn)練的控制變量實(shí)驗(yàn)。結(jié)果表明多頭訓(xùn)練可以使用較大的學(xué)習(xí)率，得到較高的準(zhǔn)確率。同時(shí)，多頭訓(xùn)練與單頭推理獲得了較高的準(zhǔn)確率，以及和基準(zhǔn)相似的推理速度。

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

下表 4 為訓(xùn)練 epoch 對(duì)結(jié)果的影響。檢測準(zhǔn)確率的增長速度在訓(xùn)練 epoch 為 144 時(shí)開始下降。因此，該研究選擇 epoch 為 144 作為基準(zhǔn)配置。

THE END