谷歌 Vision Transformers 新作：How to train your ViT?

點(diǎn)擊下方“AI算法與圖像處理”，關(guān)注一下
重磅干貨，第一時(shí)間送達(dá)

論文鏈接：https://arxiv.org/pdf/2106.10270.pdf

GitHub鏈接：

https://github.com/google-research/vision_transformer

https://github.com/rwightman/pytorch-image-models

在分享這篇文章之前，希望小伙伴能點(diǎn)點(diǎn)贊，點(diǎn)個(gè)關(guān)注，支持一下！

今年特別火的vision transformer，很多基于此的新工作，為了便于更多從業(yè)者使用ViT，這篇論文深挖一些 vision transformer 訓(xùn)練的技巧！下面一起深入看一下論文的內(nèi)容。

摘要

Vision Transformers（ViT）在圖像分類(lèi)、目標(biāo)檢測(cè)和語(yǔ)義圖像分割等領(lǐng)域具有很強(qiáng)的競(jìng)爭(zhēng)力。與卷積神經(jīng)網(wǎng)絡(luò)相比，在較小的訓(xùn)練數(shù)據(jù)集上進(jìn)行訓(xùn)練時(shí)，Vision Transformers較弱的感應(yīng)偏差通常會(huì)導(dǎo)致對(duì)模型正則化或數(shù)據(jù)增強(qiáng)（簡(jiǎn)稱(chēng)“AugReg”）的依賴(lài)性增加。為了更好地理解訓(xùn)練數(shù)據(jù)量、AugReg、模型大小和計(jì)算預(yù)算之間的相互作用，我們進(jìn)行了系統(tǒng)的實(shí)證研究。作為這項(xiàng)研究的一個(gè)結(jié)果，我們發(fā)現(xiàn)，增加計(jì)算和AugReg的組合可以產(chǎn)生與在一個(gè)數(shù)量級(jí)以上的訓(xùn)練數(shù)據(jù)上訓(xùn)練的模型具有相同性能的模型：我們?cè)诠_(kāi) 的 ImageNet-21k數(shù)據(jù)集上訓(xùn)練各種大小的ViT模型，這些模型與在更大的數(shù)據(jù)集上訓(xùn)練的對(duì)應(yīng)模型相匹配或優(yōu)于它們，但 JFT-300M 數(shù)據(jù)集尚未公開(kāi)。

Introduction

Vision Transformers（ViT）[10]最近作為卷積神經(jīng)網(wǎng)絡(luò)（CNN）的競(jìng)爭(zhēng)替代品出現(xiàn)，卷積神經(jīng)網(wǎng)絡(luò)在計(jì)算機(jī)視覺(jué)領(lǐng)域無(wú)處不在。如果沒(méi)有CNNs的平移等變性(translational equivariance)，ViT模型通常在有條件的情況下表現(xiàn)最好，大量的訓(xùn)練數(shù)據(jù)[10]或需要強(qiáng)大的AugReg方案來(lái)避免過(guò)度擬合[34]。然而，據(jù)我們所知，到目前為止，還沒(méi)有對(duì)Vision Transformers中模型正則化、數(shù)據(jù)增強(qiáng)、訓(xùn)練數(shù)據(jù)大小和計(jì)算預(yù)算之間的權(quán)衡進(jìn)行全面的研究。

在這項(xiàng)工作中，我們通過(guò)進(jìn)行深入的實(shí)證研究來(lái)填補(bǔ)這一知識(shí)空白。我們?cè)诓煌笮〉臄?shù)據(jù)集上預(yù)先訓(xùn)練了大量的ViT模型（不同大小和與resnet的混合[14]），同時(shí)對(duì)不同數(shù)量的正則化和數(shù)據(jù)擴(kuò)充進(jìn)行了精心設(shè)計(jì)的比較。我們繼續(xù)用所得到的模型進(jìn)行廣泛的遷移學(xué)習(xí)實(shí)驗(yàn)。我們主要著眼于有限的計(jì)算和數(shù)據(jù)標(biāo)注預(yù)算的情況。

研究的同質(zhì)性是本文的主要貢獻(xiàn)之一。對(duì)于絕大多數(shù)涉及 Vision Transformers 的工作來(lái)說(shuō)，在平等的基礎(chǔ)上重新訓(xùn)練所有基線和建議的方法是不實(shí)際的，特別是那些接受過(guò)大量數(shù)據(jù)訓(xùn)練的方法。此外，還有許多難以有效控制的微妙和隱含的設(shè)計(jì)選擇，例如復(fù)雜擴(kuò)充方案的精確實(shí)現(xiàn)、超參數(shù)（例如學(xué)習(xí)速率調(diào)度、權(quán)重衰減）、測(cè)試時(shí)間預(yù)處理、數(shù)據(jù)集分割等等。這種不一致可能會(huì)導(dǎo)致大量的噪聲添加到結(jié)果中，很可能影響得出任何結(jié)論的能力。因此，本工作報(bào)告所依據(jù)的所有模型都經(jīng)過(guò)了一致的訓(xùn)練和評(píng)估。

我們從研究中得出的見(jiàn)解構(gòu)成了本文的另一個(gè)重要貢獻(xiàn)。特別是，我們證明了精心選擇的正則化和增廣（從模型精度的角度）大致相當(dāng)于訓(xùn)練數(shù)據(jù)量增加了10倍。然而，不管模型是用更多的數(shù)據(jù)訓(xùn)練還是用更好的AugRegs訓(xùn)練，都必須花費(fèi)大致相同的計(jì)算量才能得到性能相似的模型。我們進(jìn)一步評(píng)估在對(duì)不同類(lèi)別的數(shù)據(jù)集微調(diào)結(jié)果模型時(shí)，添加數(shù)據(jù)或更好的AugReg之間是否存在差異。

此外，我們的目的是闡明在實(shí)踐中使用 Vision Transformers 的其他方面，例如比較中等規(guī)模數(shù)據(jù)集的遷移學(xué)習(xí)和從頭開(kāi)始的訓(xùn)練。最后，我們?cè)u(píng)估各種計(jì)算與性能的權(quán)衡。我們將在第4節(jié)中詳細(xì)討論上述所有見(jiàn)解。

實(shí)驗(yàn)設(shè)置

在這一節(jié)中，我們描述了我們的統(tǒng)一實(shí)驗(yàn)裝置，這是貫穿全文的。我們使用一個(gè)JAX/Flax[15，3]代碼庫(kù)來(lái)使用 TPU 進(jìn)行預(yù)訓(xùn)練和遷移學(xué)習(xí)。然而，使用timm PyTorch庫(kù)在V100 GPU（16G）上獲得了推斷速度測(cè)量值[37]。所有數(shù)據(jù)集都通過(guò)TensorFlow數(shù)據(jù)集庫(kù)進(jìn)行訪問(wèn)，這有助于確保一致性和再現(xiàn)性。下面提供了有關(guān)我們?cè)O(shè)置的更多詳細(xì)信息。

新發(fā)現(xiàn)

4.1 Scaling datasets with AugReg and compute

研究的一個(gè)主要發(fā)現(xiàn)(如圖1(左)所示)是，通過(guò)使用圖像增強(qiáng)和模型正則化預(yù)訓(xùn)練一個(gè)模型，使其達(dá)到與增加數(shù)據(jù)集大小約一個(gè)數(shù)量級(jí)相同的精度。更準(zhǔn)確地說(shuō)，在AugReg ImageNet-1k上訓(xùn)練的最佳模型的性能與在10倍大的普通ImageNet-21k數(shù)據(jù)集上訓(xùn)練的相同模型的性能差不多。

類(lèi)似地，在AugReg ImageNet-21k上訓(xùn)練的最佳模型，當(dāng)計(jì)算量也增加時(shí)，將匹配或優(yōu)于在普通JFT-300M數(shù)據(jù)集上訓(xùn)練的模型。因此，可以將這些結(jié)果與公開(kāi)可用的數(shù)據(jù)集進(jìn)行匹配，可以想象，在JFT-300M上進(jìn)行更長(zhǎng)時(shí)間的訓(xùn)練和使用AugReg可能會(huì)進(jìn)一步提高性能。

當(dāng)然，這些結(jié)果不能適用于任意小的數(shù)據(jù)集。只對(duì)ImageNet-1k的10%進(jìn)行大量數(shù)據(jù)增強(qiáng)的ResNet50訓(xùn)練可以改善結(jié)果，但不能恢復(fù)對(duì)完整數(shù)據(jù)集的訓(xùn)練。

4.2 Transfer is the better option

在這里，作者調(diào)查了對(duì)于從業(yè)者可能遇到的合理規(guī)模的數(shù)據(jù)集，是否建議嘗試使用AugReg從頭開(kāi)始進(jìn)行訓(xùn)練，或者是否把時(shí)間和金錢(qián)花在遷移預(yù)訓(xùn)練模型上更好。其結(jié)果是，就大多數(shù)實(shí)際目的而言，遷移預(yù)先訓(xùn)練的模型不僅成本效益更高，而且會(huì)帶來(lái)更好的結(jié)果。

作者在一個(gè)與ImageNet-1k數(shù)據(jù)集相似大小的數(shù)據(jù)集上對(duì)小的ViT-Ti/16模型進(jìn)行了搜索，尋找一個(gè)好的訓(xùn)練策略。Resisc45包含大約3萬(wàn)幅訓(xùn)練圖像，由一種非常不同的衛(wèi)星圖像組成，ImageNet-1k或ImageNet-21k都沒(méi)有很好地覆蓋這些圖像。圖1(右)和圖2顯示了這一廣泛搜索的結(jié)果。

最驚人的發(fā)現(xiàn)是，無(wú)論花費(fèi)多少訓(xùn)練時(shí)間，對(duì)于微小的Pet37數(shù)據(jù)集，似乎不可能從頭開(kāi)始訓(xùn)練ViT模型，使其達(dá)到接近遷移模型的精度。此外，由于預(yù)訓(xùn)練模型可以免費(fèi)獲取，所以從業(yè)者的預(yù)訓(xùn)練成本實(shí)際上為零，只有用于遷移學(xué)習(xí)的計(jì)算損失，因此遷移預(yù)訓(xùn)練的模型同時(shí)也大大便宜。

對(duì)于更大的Resisc45數(shù)據(jù)集，這個(gè)結(jié)果仍然成立，盡管多花費(fèi)2個(gè)數(shù)量級(jí)的計(jì)算和執(zhí)行大量搜索可能接近(但達(dá)不到)預(yù)先訓(xùn)練的模型的精度。

值得注意的是，這并沒(méi)有考慮到很難量化的“exploration cost”。對(duì)于訓(xùn)練前的模型，我們強(qiáng)調(diào)那些在訓(xùn)練前驗(yàn)證集上表現(xiàn)最好的模型，可以稱(chēng)為推薦模型?？梢钥吹?，使用推薦的模型有很高的可能性在幾次嘗試中就能獲得良好的結(jié)果。

4.3 More data yields more generic models

通過(guò)將預(yù)訓(xùn)練模型遷移到下游任務(wù)來(lái)研究預(yù)訓(xùn)練數(shù)據(jù)集大小的影響。作者在VTAB上評(píng)估了訓(xùn)練前的模型，包括19個(gè)不同的任務(wù)。

圖3顯示了3個(gè)VTAB類(lèi)別的結(jié)果:natural、specialized和structured。模型按推理時(shí)間進(jìn)行排序，模型越大推理速度越慢。

首先比較使用相同計(jì)算預(yù)算的2個(gè)模型，唯一的區(qū)別是ImageNet-1k(1.3M圖像)和ImageNet-21k (13M圖像)的數(shù)據(jù)集大小。作者實(shí)驗(yàn)對(duì)比ImageNet-1k訓(xùn)練300個(gè)epoch的模型和ImageNet-21k上訓(xùn)練30個(gè)epoch模型發(fā)現(xiàn)，在ImageNet-21k上進(jìn)行預(yù)訓(xùn)練的模型3個(gè)VTAB類(lèi)別上都明顯優(yōu)于ImageNet-1k。

隨著計(jì)算預(yù)算的不斷增長(zhǎng)，作者觀察到ImageNet-21k數(shù)據(jù)集在10倍長(zhǎng)的調(diào)度上的一致改進(jìn)。在一些幾乎已經(jīng)解決的任務(wù)中，例如花，獲得的絕對(duì)數(shù)量很小。對(duì)于剩下的任務(wù)，與短期訓(xùn)練的模型相比，改進(jìn)是顯著的。

總的來(lái)說(shuō)得出的結(jié)論是，數(shù)據(jù)越多，模型就越通用，這一趨勢(shì)適用于不同的任務(wù)。作者建議設(shè)計(jì)選擇使用更多的數(shù)據(jù)和一個(gè)固定的計(jì)算預(yù)算。

4.4 Prefer augmentation to regularization

目前尚不清楚在RandAugment和Mixup等數(shù)據(jù)增強(qiáng)和Dropout和randomdepth等模型正則化之間有哪些取舍。在本節(jié)的目標(biāo)是發(fā)現(xiàn)這些通用模式，當(dāng)將Vision transformer應(yīng)用到一個(gè)新任務(wù)時(shí)，可以作為經(jīng)驗(yàn)規(guī)則使用。

在圖4中，作者展示了為每個(gè)單獨(dú)設(shè)置獲得的上游驗(yàn)證得分，即在更改數(shù)據(jù)集時(shí)，數(shù)字是不具有可比性的。

一個(gè)單元格的顏色編碼其分?jǐn)?shù)的改善或變差，與非正則化的，未增強(qiáng)的設(shè)置，即最左邊的列。增強(qiáng)強(qiáng)度從左到右依次增大，模型容量從上到下依次增大。

第1個(gè)可見(jiàn)的觀察結(jié)果是，對(duì)于中等規(guī)模的ImageNet-1k數(shù)據(jù)集，任何類(lèi)型的AugReg都有幫助。然而，當(dāng)使用10倍大的ImageNet-21k數(shù)據(jù)集并保持計(jì)算固定時(shí)，即運(yùn)行30個(gè)epoch，任何一種AugReg都會(huì)影響除最大模型之外的所有模型的性能。只有當(dāng)計(jì)算預(yù)算增加到300個(gè)時(shí)，AugReg才幫助更多的模型，盡管即使那樣，它也繼續(xù)影響較小的模型。

一般來(lái)說(shuō)，增加增廣效果比增加正規(guī)化效果好得多。更具體地說(shuō)，圖4中每個(gè)映射右側(cè)的細(xì)列顯示，對(duì)于任何給定的模型，其最佳正則化分?jǐn)?shù)減去最佳非正則化分?jǐn)?shù)。

在圖7中，作者通過(guò)dropout和random depth的方式向模型添加正則化時(shí)，顯示了精度上的增益(綠色，正數(shù))或損失(紅色，負(fù)數(shù))。在早期的實(shí)驗(yàn)中證實(shí)，兩者結(jié)合(峰值)下降概率0.1確實(shí)是最好的設(shè)置。

這表明，模型正規(guī)化主要幫助較大的模型，但是當(dāng)訓(xùn)練時(shí)間較長(zhǎng)的情況下，特別是ImageNet-21的預(yù)訓(xùn)練，除了最大的模型它對(duì)所有的模型都有害的。

4.5 Choosing which pre-trained model to transfer

如上所述，在對(duì)ViT模型進(jìn)行預(yù)訓(xùn)練時(shí)，各種正則化和數(shù)據(jù)增強(qiáng)設(shè)置會(huì)導(dǎo)致模型具有顯著不同的性能。

然后，從實(shí)踐者的觀點(diǎn)來(lái)看，一個(gè)自然的問(wèn)題出現(xiàn)了:如何選擇一個(gè)模型進(jìn)一步適應(yīng)最終的應(yīng)用程序?

• 一種方法是：對(duì)所有可用的預(yù)訓(xùn)練模型進(jìn)行下游適應(yīng)，然后根據(jù)下游任務(wù)的驗(yàn)證分?jǐn)?shù)選擇表現(xiàn)最好的模型。但是這在實(shí)踐中可能是相當(dāng)麻煩的。

• 另一種方法是：可以根據(jù)上游驗(yàn)證精度選擇一個(gè)單獨(dú)的預(yù)訓(xùn)練模型，然后只使用該模型進(jìn)行自適應(yīng)，這要簡(jiǎn)單得多。

在這里作者將分析這2種策略之間的權(quán)衡。在5個(gè)不同的數(shù)據(jù)集上對(duì)它們進(jìn)行了大量的預(yù)訓(xùn)練模型的比較。具體來(lái)說(shuō)，在圖5(左)中強(qiáng)調(diào)了只適應(yīng)最好的預(yù)訓(xùn)練模型的簡(jiǎn)單策略和適應(yīng)所有預(yù)訓(xùn)練模型(然后選擇最好的)的復(fù)雜策略之間的性能差異。

結(jié)果好壞參半，但總體上反映出，在大多數(shù)情況下，成本較低的策略與成本較高的策略效果相同。然而，有一些顯著的異常值，當(dāng)它有利于適應(yīng)所有的模型。

因此，作者得出結(jié)論，選擇一個(gè)基于上游分?jǐn)?shù)的單一預(yù)訓(xùn)練模型是一種具有成本效益的實(shí)用策略，并在整個(gè)論文中使用它。然而，作者也強(qiáng)調(diào)，如果有額外的計(jì)算資源可用，那么在某些情況下，可以通過(guò)微調(diào)額外的預(yù)訓(xùn)練模型來(lái)進(jìn)一步提高自適應(yīng)性能。

關(guān)于ImageNet-1k數(shù)據(jù)集驗(yàn)證數(shù)據(jù)的說(shuō)明

在執(zhí)行上述分析時(shí)，作者發(fā)現(xiàn)在ImageNet-21k上預(yù)先訓(xùn)練并遷移到ImageNet-1k數(shù)據(jù)集的模型存在一個(gè)微小但嚴(yán)重的問(wèn)題。這些模型(特別是大型模型)的驗(yàn)證分?jǐn)?shù)與觀察到的測(cè)試性能沒(méi)有很好的關(guān)聯(lián)，見(jiàn)圖5(左)。這是因?yàn)镮mageNet-21k數(shù)據(jù)包含ImageNet-1k訓(xùn)練數(shù)據(jù)，作者使用訓(xùn)練數(shù)據(jù)的一個(gè)小split來(lái)進(jìn)行評(píng)估(見(jiàn)3.1節(jié))。

因此，在較長(zhǎng)訓(xùn)練計(jì)劃上的大型模型記憶了來(lái)訓(xùn)練集的數(shù)據(jù)，這使得在小評(píng)估集中計(jì)算的評(píng)估指標(biāo)存在偏差。為了解決這個(gè)問(wèn)題并支持公平的超參數(shù)選擇，作者使用獨(dú)立收集的ImageNetV2數(shù)據(jù)作為傳輸?shù)絀mageNet-1k的驗(yàn)證split。如圖5(右)所示。作者沒(méi)有在其他數(shù)據(jù)集中觀察到類(lèi)似的問(wèn)題。

作者建議將ImageNet-21k模型遷移到ImageNet-1k的研究人員遵循這一策略。

4.6 Prefer increasing patch-size to shrinking model-size

作者研究的一個(gè)意想不到的結(jié)果是，訓(xùn)練了幾個(gè)模型，它們?cè)谕评硗掏铝糠矫娲笾孪嗟?，但在質(zhì)量方面差異很大。

具體地說(shuō)，圖6(右)顯示了包含Tiny變體的模型比具有32-patch-size的類(lèi)似快速的更大模型的性能要差得多。對(duì)于給定的分辨率，patch-size會(huì)影響self-attention執(zhí)行的token數(shù)量，因此會(huì)影響模型容量，而參數(shù)計(jì)數(shù)并不能反映模型容量。參數(shù)計(jì)數(shù)既不反映速度，也不反映容量。

結(jié)論

我們首次系統(tǒng)地、大規(guī)模地研究了正則化、數(shù)據(jù)增強(qiáng)、模型大小和訓(xùn)練數(shù)據(jù)大小之間的相互作用，包括它們對(duì)達(dá)到一定性能水平所需的計(jì)算預(yù)算的影響。我們還通過(guò)遷移學(xué)習(xí)的視角來(lái)評(píng)估預(yù)先訓(xùn)練的模型。因此，我們描述了一個(gè)相當(dāng)復(fù)雜的landscape 訓(xùn)練設(shè)置的預(yù)訓(xùn)練Vision Transformers 測(cè)試不同的模型大小。我們的實(shí)驗(yàn)對(duì)各種技術(shù)的影響以及數(shù)據(jù)增強(qiáng)和正則化有益與否的情況產(chǎn)生了許多令人驚訝的見(jiàn)解。

我們還對(duì) Vision Transformers 的遷移學(xué)習(xí)環(huán)境進(jìn)行了深入的分析。我們的結(jié)論是，在廣泛的數(shù)據(jù)集中，即使下游數(shù)據(jù)的興趣似乎只是弱相關(guān)的數(shù)據(jù)用于預(yù)訓(xùn)練，遷移學(xué)習(xí)仍然是最好的選擇。我們的分析還表明，在執(zhí)行類(lèi)似的預(yù)訓(xùn)練模型，對(duì)于遷移學(xué)習(xí)，一個(gè)具有更多訓(xùn)練數(shù)據(jù)的模型應(yīng)該比一個(gè)具有更多數(shù)據(jù)擴(kuò)充的模型更為可取。

希望本文的研究能對(duì)今后Vision Transformers的研究起到一定的指導(dǎo)作用，為尋求優(yōu)化最終模型性能的從業(yè)者提供有效訓(xùn)練設(shè)置的來(lái)源，根據(jù)給定的計(jì)算預(yù)算。

參考鏈接：

https://mp.weixin.qq.com/s/9NCl_chR5QRXvphrx-IO6Q

努力分享優(yōu)質(zhì)的計(jì)算機(jī)視覺(jué)相關(guān)內(nèi)容，歡迎關(guān)注：

個(gè)人微信（如果沒(méi)有備注不拉群！）
請(qǐng)注明：地區(qū)+學(xué)校/企業(yè)+研究方向+昵稱(chēng)



下載1：何愷明頂會(huì)分享

在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù)：何愷明，即可下載。總共有6份PDF，涉及 ResNet、Mask RCNN等經(jīng)典工作的總結(jié)分析

下載2：終身受益的編程指南：Google編程風(fēng)格指南

在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù)：c++，即可下載。歷經(jīng)十年考驗(yàn)，最權(quán)威的編程規(guī)范！

下載3 CVPR2021

在「AI算法與圖像處理」公眾號(hào)后臺(tái)回復(fù)：CVPR，即可下載1467篇CVPR 2020論文 和 CVPR 2021 最新論文

點(diǎn)亮 ，告訴大家你也在看