Apple提出MobileViT:讓移動(dòng)設(shè)備用上輕量級(jí)、低延遲的視覺(jué)Transformer
在這篇論文中,來(lái)自蘋(píng)果的研究者提出了一種用于移動(dòng)設(shè)備的輕量級(jí)通用視覺(jué) transformer——MobileViT。該網(wǎng)絡(luò)在 ImageNet-1k 數(shù)據(jù)集上實(shí)現(xiàn)了 78.4% 的最佳精度,比 MobileNetv3 還要高 3.2%,而且訓(xùn)練方法簡(jiǎn)單。目前,該論文已被 ICLR 2022 接收。
- 論文鏈接:https://arxiv.org/pdf/2110.02178.pdf
- 代碼鏈接:https://github.com/apple/ml-cvnets
輕量級(jí)卷積神經(jīng)網(wǎng)絡(luò)(CNN)是移動(dòng)設(shè)備機(jī)器視覺(jué)任務(wù)的首選。它們的空間歸納偏置允許它們?cè)诓煌囊曈X(jué)任務(wù)中,可以使用較少的參數(shù)進(jìn)行視覺(jué)內(nèi)容的學(xué)習(xí)表示。但是這些網(wǎng)絡(luò)在空間上是局部的。為了學(xué)習(xí)視覺(jué)內(nèi)容的全局表示,要采用基于自注意力的視覺(jué) transformer(ViT)。與 CNN 不同,ViT 是重量級(jí)的。在這篇文章中,作者提出了以下問(wèn)題:是否有可能結(jié)合 CNN 和 ViT 的優(yōu)勢(shì),為移動(dòng)設(shè)備機(jī)器視覺(jué)任務(wù)構(gòu)建一個(gè)輕量級(jí)、低延遲的神經(jīng)網(wǎng)絡(luò)模型?
為了解決上面的問(wèn)題,作者提出了 MobileViT——一種用于移動(dòng)設(shè)備的輕量級(jí)通用視覺(jué) transformer。MobileViT 從另一個(gè)角度介紹了使用 transformer 進(jìn)行全局信息處理的方法。
具體來(lái)說(shuō),MobileViT 使用張量對(duì)局部和全局信息進(jìn)行有效地編碼(圖 1b 所示)。
與 ViT 及其變體(有卷積和無(wú)卷積)不同,MobileViT 從不同的角度學(xué)習(xí)全局表示。標(biāo)準(zhǔn)卷積涉及三個(gè)操作:展開(kāi)、局部處理和折疊。MobileViT 使用 transformer 將卷積中的局部處理方式替換為全局處理。這使得 MobileViT 兼具 CNN 和 ViT 的特性。這使它可以用更少的參數(shù)和簡(jiǎn)單的訓(xùn)練方法(如基本增強(qiáng))學(xué)習(xí)更好的表示。該研究首次證明,輕量級(jí) ViT 可以通過(guò)簡(jiǎn)單的訓(xùn)練方法在不同的移動(dòng)視覺(jué)任務(wù)中實(shí)現(xiàn)輕量級(jí) CNN 級(jí)性能。
對(duì)于大約 500-600 萬(wàn)的參數(shù),MobileViT 在 ImageNet-1k 數(shù)據(jù)集上實(shí)現(xiàn)了 78.4% 的最佳精度,比 MobileNetv3 還要高 3.2%,而且訓(xùn)練方法簡(jiǎn)單(MobileViT 與 MobileNetv3:300 vs 600 epoch;1024 vs 4096 batch size)。在高度優(yōu)化的移動(dòng)視覺(jué)任務(wù)的體系結(jié)構(gòu)中,當(dāng) MobileViT 作為特征主干網(wǎng)絡(luò)時(shí),性能顯著提高。如果將 MNASNet(Tan 等人,2019 年)替換為 MobileViT,作為 SSDLite(Sandler 等人,2018 年)的特征主干網(wǎng)絡(luò),能生成更好的(+1.8%mAP)和更小的(1.8×)檢測(cè)網(wǎng)絡(luò)(圖 2)。
架構(gòu)細(xì)節(jié)
MobileViT 塊
如圖 1b 所示的 MobileViT 塊的作用是使用包含較少參數(shù)的輸入張量學(xué)習(xí)局部和全局信息。形式上,對(duì)于給定的輸入張量 X∈ R^(H×W×C),MobileViT 應(yīng)用一個(gè) n×n 標(biāo)準(zhǔn)卷積層,然后是逐點(diǎn)(1×1)卷積層來(lái)生成 X_L∈ R^(H×W×d)。n×n 卷積層編碼局部空間信息,而逐點(diǎn)卷積通過(guò)學(xué)習(xí)輸入通道的線性組合將張量投影到高維空間(或 d 維,其中 d>C)。
有了 MobileViT,我們希望在有效感受野為 H×W 的情況下對(duì)長(zhǎng)程非局部依賴(lài)進(jìn)行建模。目前研究最多的長(zhǎng)程依賴(lài)建模方法之一是空洞卷積(dilated convolution)。然而,這種方法需要仔細(xì)選擇擴(kuò)張率(dilation rate)。否則,權(quán)重將被應(yīng)用于填充的零,而不是有效的空間區(qū)域(Yu&Koltun,2016;Chen 等人,2017;Mehta 等人,2018)。另一個(gè)候選的解決方案是自注意力(Wang 等人,2018 年;Ramachandran 等人,2019 年;Bello 等人,2019 年;Dosovitskiy 等人,2021 年)。在自注意力方法中,具有多頭自注意力的(ViT)已被證明對(duì)視覺(jué)識(shí)別任務(wù)是有效的。然而,ViT 是重量級(jí)的,并且模型優(yōu)化能力低于標(biāo)準(zhǔn),這是因?yàn)?ViT 缺少空間歸納偏置(肖等人,2021;格雷厄姆等人,2021)。
為了讓 MobileViT 學(xué)習(xí)具有空間歸納偏置的全局表示,將 X_L 展開(kāi)為 N 個(gè)非重疊 flattened patches X_U∈ R^(P×N×d)。這里,P=wh,N=HW/P 是 patch 的數(shù)量,h≤ n 和 w≤ n 分別是 patch 的高度和寬度。對(duì)于每個(gè) p∈ {1,···,P},通過(guò) transformer 對(duì) patch 間的關(guān)系進(jìn)行編碼以獲得 X_G∈ R^(P×N×d):
與丟失像素空間順序的 ViT 不同,MobileViT 既不會(huì)丟失 patch 順序,也不會(huì)丟失每個(gè) patch 內(nèi)像素的空間順序(圖 1b)。因此,我們可以折疊 X_G∈ R^(P×N×d)以獲得 X_F∈ R^(H×W×d)。然后使用逐點(diǎn)卷積將 X_F 投影到低 C 維空間,并通過(guò)級(jí)聯(lián)操作與 X 組合。然后使用另一個(gè) n×n 卷積層來(lái)融合這些連接的特征。由于 X_U(p)使用卷積對(duì) n×n 區(qū)域的局部信息進(jìn)行編碼,X_G(p)對(duì)第 p 個(gè)位置的 p 個(gè) patch 的全局信息進(jìn)行編碼,所以 X_G 中的每個(gè)像素都可以對(duì) X 中所有像素的信息進(jìn)行編碼,如圖 4 所示。因此,MobileViT 的整體有效感受野為 H×W。
與卷積的關(guān)系
標(biāo)準(zhǔn)卷積可以看作是三個(gè)連續(xù)操作:(1)展開(kāi),(2)矩陣乘法(學(xué)習(xí)局部表示)和(3)折疊。MobileViT 與卷積相似,因?yàn)樗怖昧讼嗤臉?gòu)建塊。MobileViT 用更深層的全局處理(transformer 層)取代卷積中的局部處理(矩陣乘法)。因此,MobileViT 具有類(lèi)似卷積的特性(如空間偏置)。因此,MobileViT 塊可以被視為卷積 transformer 。作者有意簡(jiǎn)單設(shè)計(jì)的優(yōu)點(diǎn)就是,卷積和 transformer 的底層高效實(shí)現(xiàn)可以開(kāi)箱即用,從而允許我們?cè)诓煌脑O(shè)備上使用 MobileViT,而無(wú)需任何額外的改動(dòng)。
輕量級(jí)
MobileViT 使用標(biāo)準(zhǔn)卷積和 transformer 分別學(xué)習(xí)局部和全局表示。相關(guān)的研究作(如 Howard et al.,2017;Mehta et al.,2021a)表明,使用這些層設(shè)計(jì)的網(wǎng)絡(luò)量級(jí)很重,因此自然會(huì)產(chǎn)生一個(gè)問(wèn)題:為什么 MobileViT 的量級(jí)很輕?作者認(rèn)為,問(wèn)題主要在于學(xué)習(xí) transformer 的全局表示。對(duì)于給定的 patch,之前的研究是(如 Touvron 等人,2021a;Graham 等人,2021)通過(guò)學(xué)習(xí)像素的線性組合將空間信息轉(zhuǎn)換為潛在信息(圖 1a)。然后,通過(guò)使用 transformer 學(xué)習(xí) patch 間的信息,對(duì)全局信息進(jìn)行編碼。因此,這些模型失去了 CNN 固有的圖像特定歸納偏置。因此,它們需要更強(qiáng)的能力來(lái)學(xué)習(xí)視覺(jué)表示。這就導(dǎo)致這些網(wǎng)絡(luò)模型既深又寬。與這些模型不同,MobileViT 使用卷積和 transformer 的方式是,生成的 MobileViT 既具有類(lèi)似卷積的屬性,又同時(shí)允許全局處理。這種建模能力使我們能夠設(shè)計(jì)淺層和窄層的 MobileViT 模型,因此最終的模型很輕。與使用 L=12 和 d=192 的基于 ViT 的模型 DeIT 相比,MobileViT 模型分別在大小為 32×32、16×16 和 8×8 的空間層次上使用 L={2,4,3}和 d={96,120,144},產(chǎn)生的 MobileViT 網(wǎng)絡(luò)比 DeIT 網(wǎng)絡(luò)更快(1.85×)、更?。?×)、更好(+1.8%)(表 3 所示)。
計(jì)算成本
MobileViT 和 ViTs(圖 1a)中多頭自注意力的計(jì)算成本分別為 O(N^2Pd)和 O(N^2d)。理論上 MobileViT 效率是比 ViTs 低的。然而在實(shí)踐中,MobileViT 實(shí)際比 ViTs 更高效。在 ImageNet-1K 數(shù)據(jù)集上,與 DeIT 相比,MobileViT 的 FLOPs 減少了一半,并且精確度提高了 1.8%(表 3 所示)。這是因?yàn)檩p量級(jí)設(shè)計(jì)(前面討論)的原因。
MobileViT 架構(gòu)
作者設(shè)計(jì)的網(wǎng)絡(luò)也是受到輕量級(jí) CNN 理念的啟發(fā)。以三種不同的網(wǎng)絡(luò)大?。⊿:small、XS:extra-small 和 XXS:extra-extra-small)訓(xùn)練 MobileViT 模型,這些網(wǎng)絡(luò)通常用于移動(dòng)視覺(jué)任務(wù)(圖 3c)。MobileViT 中的初始層是一個(gè) 3×3 的標(biāo)準(zhǔn)卷積,然后是 MobileNetv2(或 MV2)塊和 MobileViT 塊(圖 1b 和 §A)。使用 Swish(Elfwing 等人,2018)作為激活函數(shù)。按照 CNN 模型,在 MobileViT 塊中使用 n=3。特征映射的空間維度通常是 2 和 h、w 的倍數(shù)≤ n。因此在所有空間級(jí)別設(shè)置 h=w=2。MobileViT 網(wǎng)絡(luò)中的 MV2 模塊主要負(fù)責(zé)下采樣。因此,這些區(qū)塊在 MobileViT 網(wǎng)絡(luò)中是淺而窄的。圖 3d 中 MobileViT 的空間水平參數(shù)分布進(jìn)一步表明,在不同的網(wǎng)絡(luò)配置中,MV2 塊對(duì)網(wǎng)絡(luò)參數(shù)的貢獻(xiàn)非常小.
實(shí)驗(yàn)結(jié)果
IMAGENET-1K 數(shù)據(jù)集上的圖像分類(lèi)結(jié)果
與 CNN 對(duì)比
圖 6a 顯示,在不同的網(wǎng)絡(luò)規(guī)模(MobileNet v1(Howard et al.,2017)、MobileNet v2(Sandler et al.,2018)、ShuffleNet v2(Ma et al.,2018)、ESPNetv2(Mehta et al.,2019)和 MobileNet v3(Howard et al.,2019))中,MobileNet 在性能上優(yōu)于輕量級(jí) CNN。對(duì)于大約 250 萬(wàn)個(gè)參數(shù)的模型(圖 6b),在 ImageNet1k 驗(yàn)證集上,MobileViT 的性能比 MobileNetv2 好 5%,比 ShuffleNetv2 好 5.4%,比 MobileNetv3 好 7.4%。圖 6c 進(jìn)一步表明,MobileViT 的性能優(yōu)于重量級(jí) CNN(ResNet(He 等人,2016 年)、DenseNet(Huang 等人,2017 年)、ResNet SE(Hu 等人,2018 年)和 EfficientNet(Tan&Le,2019a))。對(duì)于類(lèi)似數(shù)量的參數(shù),MobileViT 比 EfficientNet 的準(zhǔn)確度高 2.1%。
與 ViTs 進(jìn)行比較
圖 7 將 MobileViT 與在 ImageNet-1k 未蒸餾數(shù)據(jù)集上從頭開(kāi)始訓(xùn)練的 ViT 變體進(jìn)行了比較(DeIT(Touvron et al.,2021a)、T2T(Yuan et al.,2021b)、PVT(Wang et al.,2021)、CAIT(Touvron et al.,2021b)、DeepViT(Zhou et al.,2021)、CeiT(Yuan et al.,2021a)、CrossViT(Chen et al.,2021a)、LocalViT(Li et al.,2021)、PiT(Heo et al.,2021),ConViT(d’Ascoli 等人,2021 年)、ViL(Zhang 等人,2021 年)、BoTNet(Srinivas 等人,2021 年)和 Mobile-former(Chen 等人,2021b 年)。不像 ViT 變體顯著受益于深層數(shù)據(jù)增強(qiáng)(例如,PiT w / 基礎(chǔ)與高級(jí):72.4(R4)與 78.1(R17);圖 7b),MobileViT 通過(guò)更少的參數(shù)和基本的增強(qiáng)實(shí)現(xiàn)了更好的性能。例如,MobileViT 只有 DeIT 的 1/2.5 大小,但性能比 DeIT 好 2.6%(圖 7b 中的 R3 和 R8)。
MOBILEVIT 作為通用主干網(wǎng)絡(luò)的表現(xiàn)
移動(dòng)目標(biāo)檢測(cè)
表 1a 顯示,對(duì)于相同的輸入分辨率 320×320,使用 MobileViT 的 SSDLite 優(yōu)于使用其他輕量級(jí) CNN 模型(MobileNetv1/v2/v3、MNASNet 和 MixNet)的 SSDLite。此外,使用 MobileViT 的 SSDLite 性能優(yōu)于使用重型主干網(wǎng)絡(luò)的標(biāo)準(zhǔn) SSD-300,同時(shí)學(xué)習(xí)的參數(shù)也明顯減少(表 1b)。
移動(dòng)語(yǔ)義分割:從表 2 可見(jiàn),使用 MobileViT 的特征主干網(wǎng)絡(luò)比 DeepLabv3 更小、更好。
移動(dòng)設(shè)備上的性能測(cè)試
輕量級(jí)和低延遲的網(wǎng)絡(luò)對(duì)于實(shí)現(xiàn)移動(dòng)視覺(jué)應(yīng)用非常重要。為了證明 MobileViT 對(duì)此類(lèi)應(yīng)用的有效性,使用公開(kāi)的 CoreMLTools(2021 年)將預(yù)先訓(xùn)練的全精度 MobileViT 模型轉(zhuǎn)換為 CoreML。然后在移動(dòng)設(shè)備 (iPhone12) 上測(cè)試它們的推理時(shí)間(平均超 100 次迭代)。
圖 8 顯示了 MobileViT 網(wǎng)絡(luò)在三個(gè)不同任務(wù)上的推斷時(shí)間,其中兩個(gè) patch 大小設(shè)置(Config-A: 2, 2, 2 和 Config-B: 8, 4, 2)。
? ?