開源 BiT:計算機視覺大規模預訓練的探索

文 / Lucas Beyer 和 Alexander Kolesnikov,研究工程師,Google Research

計算機視覺領域的研究者普遍認爲,現代深度神經網絡的發展始終需要使用大量標記數據,如,當前最前沿 (SOTA) 的卷積神經網絡 (CNN) 便需要在 OpenImages 或 Places(其中包含超過 100 萬張標籤圖像)等數據集上進行訓練。但是,對很多普通的應用開發者來說,收集到足夠多的標記數據頗爲困難。

爲緩解計算機視覺任務缺乏標記數據的情況,我們通常會使用由通用數據(如 ImageNet) 預訓練 的模型,即將通用數據上習得的視覺特徵在興趣任務上重複使用。儘管這種預訓練模型在實踐中相當有效,但在多環境中快速掌握和新概念理解上,仍存在短板。鑑於 BERT 和 T5 在語言領域取得的進展,我們認爲 大規模 預訓練同樣能提升計算機視覺模型的性能。

  • T5
    https://ai.googleblog.com/2020/02/exploring-transfer-learning-with-t5.html

在《Big Transfer (BiT):通用視覺表徵學習(Big Transfer (BiT): General Visual Representation Learning) 一文中,我們設計出一種新方法,利用超出業界標準規模 (ILSVRC-2012) 的圖像數據集對一般特徵進行有效的預訓練。值得一提的是,隨着預訓練數據量的增加,選擇合適的歸一化層和適當擴展架構容量十分重要。

  • Big Transfer (BiT):通用視覺表徵學習
    https://arxiv.org/abs/1912.11370

新方法表現出前所未有的性能,可適應各種新的視覺任務,包括小樣本識別設置和最近引入的“現實世界 (ObjectNet)” 基準。我們很高興能分享在公開數據集上預訓練的最佳 BiT 模型,以及 TF2、Jax 和 PyTorch 中的代碼。藉助這些內容,任何人都能在他們感興趣的任務中獲得出色性能,即使每個類只有少量標籤圖像也是如此。

  • 最佳 BiT 模型
    https://tfhub.dev/google/collections/bit/1

預訓練

爲了研究數據規模對模型的影響,我們通過三個數據集重新審視了預訓練設置中常用的設計選項(如激活函數和權重的歸一化、模型寬度/深度和訓練計劃):ILSVRC-2012(128 萬張圖像,1000 個類)、ImageNet-21k(1400 萬張圖像,約 21000 個類)以及 JFT(3 億張圖像,大約 18000 個類)。更重要的是,利用這些數據集,我們能夠專心研究先前未曾充分探索的大型數據模式。

我們首先研究了數據集規模與模型容量之間的相互作用。爲此,我們選擇了經典的 ResNet 架構進行訓練。該架構性能良好,同時簡單且可重現。從標準 50 層深的“R50x1”到 4 倍寬度、152 層深的“R152x4”的變體模型,我們在上述每個數據集上都進行了訓練。我們通過觀察發現了一個重要結論:要充分利用大量數據的優勢,就必須增加模型容量。下圖左側面板中的紅色箭頭便證明了這一結論:

左側:想要有效利用大型數據集進行預訓練,就需要增加模型容量。紅色箭頭證明了以下內容:在較大的 ImageNet-21k 上預訓練後,小型架構(較小的點)的性能會變差,而大型架構(較大的點)的性能則有所改善。右側:僅利用較大數據集進行預訓練未必會提高性能,例如,從 ILSVRC-2012 改用相對較大的 ImageNet-21k 後,性能沒提升。不過,通過增加計算預算和訓練時間,性能明顯有所提升

經觀察後,我們得出第二個更爲重要的結論:訓練時間的作用不容忽視。如果不調整計算預算、不增加訓練時間,而只增加預訓練所用數據集的規模,模型性能可能會變得更差。但是,調整新數據集的訓練時間後,性能得到了顯著提升。

我們在探索期間發現另一項修改對於性能提升也十分關鍵。我們由此得出結論,將批標準化(Batch Normalization,BN,一種可通過標準化激活函數實現穩定訓練的通用層)替換爲組標準化 (Group Normalization,GN) 對於大規模預訓練十分有利。

首先,BN 的狀態(神經激活函數的均值和方差)需要在預訓練和遷移之間進行調整,而 GN 是無狀態的,從而避免了這一難題。其次,BN 使用批處理級別的統計數據,對於每個設備上的小型批次大小來說並不可靠,而大型模型無法避免在多設備上進行訓練。由於 GN 不需要計算批處理級別的統計信息,因此也避免了這個問題。

請參閱我們的論文,以便了解更多技術細節,包括如何使用權重標準化 (Weight Standardization) 技術來確保行爲的穩定。

  • 權重標準化
    https://arxiv.org/abs/1903.10520

預訓練策略彙總:使用標準的 ResNet,增加深度和寬度,將 BatchNorm (BN) 替換爲 GroupNorm 和權重標準化 (GNWS),然後通過一個大型通用數據集進行訓練,並完成多次迭代

遷移學習

根據 BERT 在語言領域建立的方法,我們利用來自感興趣的各項“下游”任務的數據對預訓練的 BiT 模型進行了微調,這些任務可能只有少量標記數據。由於預訓練模型對視覺世界的理解能力較高,因此這個簡單的策略非常有效。

微調時需要選擇大量超參數,如學習率、權重衰減等。我們提出了一個啓發式的超參數選擇方法,名爲 “BiT-HyperRule”。該方法只以高層次的數據集特徵爲基礎,如圖像分辨率和標籤樣本數量。我們已成功將 BiT-HyperRule 方法應用到 20 多項從自然圖像到醫學影像等不同的任務上。

經預訓練後,BiT 模型可根據任何任務進行微調,即使提供的標籤樣本非常有限也無妨

將 BiT 遷移到僅有少量示例的任務後,我們發現:如果同時增加用於預訓練的通用數據量和架構容量,生成的模型明顯能更快適應新數據。在擁有 1 個樣本和 5 個樣本的 CIFAR(見下圖)上,如果在 ILSVRC 上進行預訓練,則增加模型容量所帶來的增益比較有限(綠色曲線)。但是,如果在 JFT 上進行大規模預訓練,每次增加模型容量都會帶來巨大增益(棕色曲線)。BiT-L 在 1 個樣本上的準確率可達 64%,在 5 個樣本上的準確率可達 95%。

曲線描繪了超過 5 次獨立運行(光點)後,模型遷移到 CIFAR-10(每個類有 1 張或 5 張圖像,共 10 張或 50 張圖像)上的平均準確率。事實證明,在大型數據集上進行預訓練的大型架構更能發揮數據的作用

爲了驗證該結果具有普遍性,我們還在 VTAB-1k 上對 BiT 進行了評估。VTAB-1k 是由 19 個不同任務組成的任務包,每個任務只有 1000 個標籤樣本。我們將 BiT-L 模型遷移到所有這些任務上,總體得分爲 76.3%,比之前的最高水準高出了 5.8%。

我們根據幾個標準計算機視覺基準(如 Oxford Pets、Flowers、CIFAR 等)對 BiT-L 進行了評估,並由此證明,大規模預訓練和簡單遷移的策略十分有效,即使提供的數據量僅屬於中等水平也是如此。綜上所述,BiT-L 已達到或超越了目前的最高水準。最後,我們在 MSCOCO-2017 檢測任務中將 BiT 用作 RetinaNet 的主幹,同時證實,即使是處理這種結構化輸出任務,使用大規模預訓練也大有裨益。

左側:在各種標準計算機視覺基準,BiT-L 與先前最前沿的通用模型的準確率。右側:在 MSCOCO-2017 中將 BiT 用作 RetinaNet 主幹所得出的平均精度 (AP) 結果

需要強調的是,在我們考慮的所有不同下游任務中,我們不會對任務逐個進行超參數調整,而是藉助 BiT-HyperRule。如我們的論文所述,在足夠大的驗證數據上調整超參數可以獲得更好的結果。

在 ObjectNet 上的評估

爲進一步評估 BiT 在更具挑戰的場景中的穩健性,我們在最近引入的 ObjectNet 數據集上對經 ILSVRC-2012 微調過的 BiT 模型進行了評估。ObjectNet 更爲接近真實場景,其中,對象可能出現在非典型的上下文中,並以非典型的視角和旋轉呈現。有趣的是,受益於更大的數據和架構容量,BiT-L 在這個任務上取得了令人矚目的成就:實現了高達 80.0% 的 top-5 準確率,比之前的最高水準 (SOTA) 高出了近 25%。

  • ObjectNet
    http://objectnet.dev/

BiT 在 ObjectNet 數據集上得出的評估結果。左側:top-5 準確率,右側:top-1 準確率

結論

我們發現,在利用大量通用數據進行預訓練的情況下,一個簡單的遷移策略就會讓結果大幅改觀,不管是對大型數據集,還是對數據極少、甚至每個類僅有一張圖像的任務來說均是如此。

我們已經發布 BiT-M 模型(經 ImageNet-21k 預訓練的 R152x4 模型),同時還提供了在 Jax、TensorFlow 2 和 PyTorch 中進行遷移的 Colab。除了發佈代碼外,我們還向讀者介紹了有關如何使用 BiT 模型的 TensorFlow 2 動手教程。我們期待從業者和研究人員可以從中獲益,並用該模型替代常用的 ImageNet 預訓練模型。

  • 已經發布 BiT-M 模型
    https://github.com/google-research/big_transfer

致謝

我們要感謝 Xiaohua Zhai、Joan Puigcerver、Jessica Yung、Sylvain Gelly 和 Neil Houlsby,他們共同撰寫了 BiT 論文,並參與了 BiT 開發的各項環節,同時也要感謝蘇黎世的 Brain 團隊。我們還要感謝 Andrei Giurgiu 在調試輸入流水線方面給予的幫助。感謝 Tom Small,是他製作了這篇博文中的動畫。最後,感興趣的讀者可以瞭解一下由 Google Research 的同事開發的一些與本方向相關的方法、Noisy Student 模型以及 Facebook Research 發佈的一篇高度相關的論文探索弱監督預訓練的限制 (Exploring the Limits of Weakly Supervised Pretraining)。

  • Noisy Student 模型
    https://arxiv.org/abs/1911.04252

  • 探索弱監督預訓練的限制
    https://arxiv.org/abs/1805.00932

更多 AI 相關閱讀:

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章