NeurIPS 2022 | 四分鐘內就能訓練目標檢測器，商湯基模型團隊是怎么做到的？

來自商湯的基模型團隊和香港大學等機構的研究人員提出了一種大批量訓練算法 AGVM，該研究已被NeurIPS 2022接收。

本文提出了一種大批量訓練算法 AGVM (Adaptive Gradient Variance Modulator)，不僅可以適配于目標檢測任務，同時也可以適配各類分割任務。AGVM 可以把目標檢測的訓練批量大小擴大到 1536，幫助研究人員四分鐘訓練 Faster R-CNN，3.5 小時把 COCO 刷到 62.2 mAP，均打破了目標檢測訓練速度的世界紀錄。

• 論文地址：https://arxiv.org/pdf/2210.11078.pdf

• 代碼地址：https://github.com/Sense-X/AGVM

  
  

在當前的機器學習社區中，有三個普遍的趨勢。首先，神經網絡模型會越來越大。在 NLP 領域中最大規模的模型已經達到了上萬億級別。在視覺領域，最大規模的模型也達到了三百億的量級。其次，訓練的數據集也變得越來越大。比如，ImageNet 21k 和谷歌的 JFT 數據集都具有相當規模的數據集。另外，由于數據集變得越來越大，訓練 SOTA 模型的開銷越來越大。

因此，提升訓練效率就變得愈發重要。而分布式訓練因為其適應于數據并行、模型并行和流水線并行的加速訓練方法的同時，也具備較高的 Deep Learning 通信效率而被廣泛認為是一個有效的解決方案。

隨著大模型時代的到來，目標檢測器的訓練速度越來越成為學術界和工業界的瓶頸，例如，在 COCO 的標準 setting 上把 mAP 訓到 62 以上大概需要三天的時間，算上調試成本，這在業界幾乎是不可接受的。那么，我們能不能把這個訓練時間壓到小時級別呢？事實上，在圖片分類和自然語言處理任務上，先前的研究人員借助 32K 的批量大小(batch size)，只需 14 分鐘就可以完成 ImageNet 的訓練，76 分鐘完成 Bert 的訓練。但是，在目標檢測領域，還很欠缺這類研究，導致研究人員無法充分利用當前的算力，數據集和大模型。

大批量訓練算法 AGVM 便是這個問題的最佳解決方案之一。為了支持如此大批量的訓練，同時保持模型的訓練精度，本研究提出了一套全新的訓練算法，根據密集預測不同模塊的梯度方差(gradient variance)，動態調整每一個模塊的學習率。作者在大量的密集預測網絡和數據集上進行了實驗，并且證實了該方法的合理性。

方法介紹

大批量訓練是加速大型分布式系統中深度神經網絡訓練的關鍵。尤其是在如今的大模型時代，如果不采用大批量訓練，一個網絡的訓練時間幾乎是難以接受的。但是，大批量訓練很難，因為它會產生泛化差距(generalization gap)， 直接訓練會導致其準確率降低。此前的大批量工作往往針對于圖像分類以及一些自然語言處理的任務，但密集預測任務（包括檢測分割等），同樣在視覺中處于舉足輕重的位置，此前的方法并不能在密集預測任務上有很好的表現，甚至結果比基準線更差，這導致我們難以快速訓練一個目標檢測器。

為了解決這個問題，研究人員進行了大量的實驗。最后發現，相較于傳統的分類網絡，利用密集預測網絡一個很重要的特征：密集預測網絡往往是由多個組件組成的，以 Faster R-CNN 為例：它由四個部分組成，骨干網絡 (Backbone)，特征金字塔網絡(FPN)，區域生成網絡(RPN) 和檢測頭網絡(head)，我們可以發現一個很有效的指標：密集預測網絡不同組件的梯度方差，在訓練批量很小時（例如 32），幾乎是相同的，但當訓練批量很大時（例如 512），它們呈現出很大的區別，如下圖所示：

那么，能不能直接把這些拉平呢？這直接引出了 AGVM 算法。以隨機梯度下降算法為例，上角標 i 代表第 i 個網絡模塊（例如 FPN 等），上角標 1 代表骨干網絡，代表學習率，錨定骨干網絡，可以直接將不同網絡組件的梯度 g 的方差：

梯度的方差可以由以下式子估計：

方差的具體求解細節可以參考原文，本研究同樣引入了滑動平均機制，防止網絡訓練發散。同時，研究證明了 AGVM 在非凸情況下的收斂性，討論了動量以及衰減的處理方式，具體實現細節可以參考原文。

實驗過程

本研究首先在目標檢測、實例分割、全景分割和語義分割的各種密集預測網絡上進行了測試，通過下表可以看到，當用標準批量大小訓練時，AGVM 相較傳統方法沒有明顯優勢，但當在超大批量下訓練時，AGVM 相較傳統方法擁有壓倒性的優勢，下圖第二列從左至右分別表示目標檢測，實例分割，全景分割和語義分割的表現，AGVM 超越了有史以來的所有方法：

下表詳細對比了 AGVM 和傳統方法，體現出了本研究方法的優勢：

同時，為了說明 AGVM 的優越性，本研究進行了以下三個超大規模的實驗。研究人員把 Faster R-CNN 的 batch size 放到了 1536，這樣利用 768 張 A100 可以在 4.2 分鐘內完成訓練。其次，借助 UniNet-G，本研究可以在利用 480 張 A100 的情況下，3.5 個小時讓模型在 COCO 上達到 62.2mAP（不包括骨干網絡預訓練的時間），極大的減小了訓練時間：

甚至，在 RetinaNet 上，本研究把批量大小擴展到 10K。這在目標檢測領域是從未見的批量大小，在如此大的批量下，每一個 epoch 只有十幾個迭代次數，AGVM 在如此大的批量下，仍然能展現出很強的穩定性，性能如下圖所示：

結果分析

本研究探究了一個很重要的問題：以 RetinaNet 為例，如下圖第一列所示，探究為什么會出現梯度方差不匹配這一現象。

本研究認為，這一現象來自于：網絡不同模塊間的有效批量大小 (effective batch size) 是不同的。例如，RetinaNet 的頭網絡的輸入是由特征金字塔的五層網絡輸出的，特征金字塔的 top-down 和 bottom-up pathways，以及像素維度的損失函數計算會導致頭網絡和骨干網絡的等效批量大小不同，這一原理導致了梯度方差不匹配的現象。

為了驗證這一假設，本研究依次給每一層特征使用單獨的頭網絡，移去特征金字塔網絡，隨機忽略掉 75% 的用于計算損失函數的像素，最終，本研究發現骨干網絡和頭網絡的梯度方差曲線重合了，本研究也對 Faster R-CNN 做了類似的實驗，如下圖第二列所示，更多的討論請參見原文。