Apple提出FastViT：快速卷積和Transformer混合架構(gòu)

導(dǎo)讀

本文提出了一種通用的 CNN 和 Transformer 混合的視覺(jué)基礎(chǔ)模型，移動(dòng)設(shè)備和 ImageNet 數(shù)據(jù)集上的精度相同的前提下，F(xiàn)astViT 比 CMT 快3.5倍，比 EfficientNet 快4.9倍，比 ConvNeXt 快1.9倍

太長(zhǎng)不看版

本文是 MobileOne 原班人馬打造，可以看做是 MobileOne 的方法在 Transformer 上的一個(gè)改進(jìn)型的應(yīng)用。作者取名 FastViT，是一種 CNN，Transformer 混合架構(gòu)的低延時(shí)模型。作者引入了一種新的 token mixer，叫做 RepMixer，它使用結(jié)構(gòu)重新參數(shù)化技術(shù)，通過(guò)刪除網(wǎng)絡(luò)中的 Shortcut 來(lái)降低內(nèi)存訪(fǎng)問(wèn)成本。

進(jìn)一步使用大核卷積使得 FastViT 精度得到提升，而且不怎么影響延時(shí)。在移動(dòng)設(shè)備和 ImageNet 數(shù)據(jù)集上的精度相同的前提下，F(xiàn)astViT 比 CMT 快3.5倍，比 EfficientNet 快4.9倍，比 ConvNeXt 快1.9倍。在類(lèi)似的延遲下，F(xiàn)astViT 在 ImageNet 上獲得的 Top-1 準(zhǔn)確率比 MobileOne 高 4.2%，是一種極具競(jìng)爭(zhēng)力的混合架構(gòu)模型。

1 FastViT：快速卷積 Transformer 的混合視覺(jué)架構(gòu)

論文名稱(chēng)：FastViT: A Fast Hybrid Vision Transformer using Structural Reparameterization

在CVer微信公眾號(hào)后臺(tái)回復(fù)：FastViT，可以下載本論文pdf、代碼

論文：https://arxiv.org/pdf/2303.14189

代碼：https://github.com/apple/ml-fastvit

1.1 背景和動(dòng)機(jī)

本文的目標(biāo)是做一個(gè)卷積，Attention 的低延時(shí)混合架構(gòu)，因?yàn)檫@種架構(gòu)有效地結(jié)合了 CNN 和 Transformer 的優(yōu)勢(shì)，在多個(gè)視覺(jué)任務(wù)上有競(jìng)爭(zhēng)力。本文的目標(biāo)是建立一個(gè)模型，實(shí)現(xiàn) SOTA 的精度-延時(shí) Trade-off。

本文的出發(fā)點(diǎn)是最近的像 CMT[1]，LIT[2]等 CNN 和 Transformer 混合架構(gòu)的模型都遵循 MetaFormer[3] 的架構(gòu)，它由帶有 skip-connection 的 token mixer 和帶有 skip-connection 的前饋網(wǎng)絡(luò) (Feed Forward Network) 組成。由于增加了內(nèi)存訪(fǎng)問(wèn)成本 (memory access cost)，這些跳過(guò)連接在延遲方面占了很大的開(kāi)銷(xiāo)。為了解決這個(gè)延遲開(kāi)銷(xiāo)，本文提出 RepMixer，這是一個(gè)完全可以重參數(shù)化的令牌混合器，它的特點(diǎn)1是使用結(jié)構(gòu)重參數(shù)化來(lái)刪除 skip-connection。

RepMixer 的特點(diǎn)2是在訓(xùn)練期間為主要的層添加一些過(guò)參數(shù)化的額外的分支，以在訓(xùn)練時(shí)提升模型的精度，在推理時(shí)全部消除。RepMixer 的特點(diǎn)3是在網(wǎng)絡(luò)中使用了大核卷積在前幾個(gè)階段替換掉 Self-Attention。具體是在前饋網(wǎng)絡(luò) (FFN) 層和 Patch Embedding 層中加入了大核卷積。這些更改對(duì)模型的總體延遲影響很小，同時(shí)提高了性能。

對(duì)于性能這塊作者在 iPhone 12 Pro 設(shè)備和 NVIDIA RTX-2080Ti desktop GPU 上進(jìn)行了詳盡的分析，實(shí)驗(yàn)結(jié)果如下圖1所示?？梢钥吹皆趦煞N設(shè)備上，F(xiàn)astViT 都實(shí)現(xiàn)了最佳的精度-延時(shí)的權(quán)衡。

圖1：iPhone 12 Pro 設(shè)備和 NVIDIA RTX-2080Ti desktop GPU 上的精度-延時(shí)比較

1.2 FastViT 模型架構(gòu)

FastViT 整體架構(gòu)如下圖2所示。

Stage 的內(nèi)部架構(gòu)

FastViT 采用了4個(gè) stage 的架構(gòu)，每個(gè) stage 相對(duì)于前一個(gè)的分辨率減半，通道數(shù)加倍。前3個(gè) stage 的內(nèi)部架構(gòu)是一樣的，都是訓(xùn)練的時(shí)候采用下式：

推理的時(shí)候采用結(jié)構(gòu)重參數(shù)化得到下式：

第4個(gè) stage 的內(nèi)部架構(gòu)如圖2 (a) 所示，采用 Attention 來(lái)作為 token mixer，可能是為了性能考慮，寧愿不采用結(jié)構(gòu)重參數(shù)化，犧牲延時(shí)成本，以換取更好的性能。

值得注意的是，每個(gè) Stage 中的 FFN 使用的并不是傳統(tǒng)的 FFN 架構(gòu)，而是如圖2 (c) 所示的，帶有大核 7×7 卷積的 ConvFFN 架構(gòu)。

圖2：FastViT 模型架構(gòu)

Stem 的結(jié)構(gòu)

Stem 是整個(gè)模型的起點(diǎn)，如圖2 (b) 所示，F(xiàn)astViT 的 Stem 在推理時(shí)的結(jié)構(gòu)是 3×3 卷積 + 3×3 Depth-wise 卷積 + 1×1 卷積。在訓(xùn)練時(shí)分別加上 1×1 分支或者 Identity 分支做結(jié)構(gòu)重參數(shù)化。

Patch Embedding 的架構(gòu)

Patch Embedding 是模型在 Stage 之間過(guò)渡的部分，F(xiàn)astViT 的 Patch Embedding 如圖2 (a) 所示，在推理時(shí)的結(jié)構(gòu)是 7×7 大 Kernel 的 Depth-wise 卷積 + 1×1 卷積。在訓(xùn)練時(shí)分別加上 3×3 分支做結(jié)構(gòu)重參數(shù)化。

位置編碼

位置編碼使用條件位置編碼，它是動(dòng)態(tài)生成的，并以輸入 token 的局部鄰域?yàn)闂l件。這些編碼是由 depth-wise 運(yùn)算符生成的，并添加到 Patch Embedding 中。

1.3 RepMixer 的延時(shí)優(yōu)勢(shì)

如下圖3所示，作者對(duì)比了 RepMixer 和高效的 Pooling 操作的延時(shí)情況。架構(gòu)使用的是 MetaFormer S12，大概有 1.8 GFLOPs。作者在 iPhone 12 Pro 移動(dòng)設(shè)備上為從 224×224 到 1024×1024 的各種輸入分辨率的模型計(jì)時(shí)。可以看到 RepMixer 明顯優(yōu)于 Pooling，尤其是在較高分辨率的時(shí)候。在分辨率為 384×384 時(shí)，使用 RepMixer 可以降低 25.1% 的延遲，而在分辨率為 1024×1024 時(shí)，使用 RepMixer 可以降低 43.9% 的延遲。

圖3：RepMixer 的延時(shí)優(yōu)勢(shì)

1.4 FastViT 的大核卷積

RepMixer 的感受野是局部的。我們知道 Self-Attention 操作的感受野是全局的，但是 Self-Attention 操作計(jì)算量昂貴，因此之前有工作說(shuō)使用大核卷積可以在計(jì)算量得到控制的情況下有效增加感受野的大小。FastViT 在兩個(gè)位置引入了大核卷積，分別是 Patch Embedding 層和 FFN。對(duì)比實(shí)驗(yàn)的結(jié)果如下圖4所示。將 V5 與 V3 進(jìn)行比較，模型大小增加了 11.2%，延遲增加了 2.3 倍，而 Top-1 精度的增益相對(duì)較小，只有 0.4%，說(shuō)明使用大核卷積來(lái)替換 Self-Attention 是一種高效，節(jié)約延時(shí)的方式。V2 比 V4 大 20%，延時(shí)比 V4 高 7.1%，同時(shí)在 ImageNet 上獲得相似的 Top-1 精度。

圖4：大核卷積的消融實(shí)驗(yàn)

隨著感受野的增加，大核卷積也有助于提高模型的魯棒性。FastViT 各種模型的超參數(shù)配置如下圖5所示。

圖5：FastViT 的超參數(shù)配置

1.5 實(shí)驗(yàn)結(jié)果

ImageNet-1K 圖像分類(lèi)實(shí)驗(yàn)結(jié)果

如下圖6所示是 ImageNet-1K 圖像分類(lèi)實(shí)驗(yàn)結(jié)果。對(duì)于 iPhone 設(shè)備延時(shí)的測(cè)量，作者使用 Core ML Tools (v6.0) 導(dǎo)出模型，并在帶有 iOS 16 的 iPhone12 Pro Max 上運(yùn)行，并將所有模型的 Batch Size 大小設(shè)置為1。對(duì)于 GPU延時(shí)的測(cè)量，作者把模型導(dǎo)出為 TensorRT (v8.0.1.6) 格式，并在 NVIDIA RTX-2080Ti 上運(yùn)行，Batch Size 大小為8，報(bào)告100次運(yùn)行的中位數(shù)。

與 SOTA 模型的性能比較如下圖6所示。本文的 FastViT 實(shí)現(xiàn)了最佳的精度-延時(shí)均衡，比如 FastViT-S12 在 iPhone 12 Pro 上比 MobileOne-S4 快 26.3%，GPU 上快 26.9%。在 83.9% 的 Top-1 精度下，F(xiàn)astViT-MA36 比 iPhone 12 Pro 上優(yōu)化的 ConvNeXt-B 模型快 1.9倍, GPU上快2.0倍。

圖6：ImageNet-1K 圖像分類(lèi)實(shí)驗(yàn)結(jié)果

知識(shí)蒸餾實(shí)驗(yàn)結(jié)果

如下圖7所示是 FastViT 作為學(xué)生模型的知識(shí)蒸餾實(shí)驗(yàn)結(jié)果。作者遵循 DeiT 中的實(shí)驗(yàn)設(shè)置，RegNet16GF 作為教師模型，使用 Hard Distillation，其中教師的輸出設(shè)置為 true label，一共訓(xùn)練300個(gè) Epochs。FastViT 優(yōu)于最近最先進(jìn)的模型 EfficientFormer。FastViT-SA24 的性能與 EfficientFormer-L7 相似，但參數(shù)少3.8倍，F(xiàn)LOPs 少2.7倍，延遲低2.7倍。

圖7：知識(shí)蒸餾實(shí)驗(yàn)結(jié)果

目標(biāo)檢測(cè)和語(yǔ)義分割實(shí)驗(yàn)結(jié)果

對(duì)于語(yǔ)義分割，作者在 ADE20k 上驗(yàn)證了模型的性能語(yǔ)義分割模型頭使用的是 Semantic FPN，所有的模型都是用預(yù)先訓(xùn)練好的對(duì)應(yīng)圖像分類(lèi)模型的權(quán)重進(jìn)行初始化。在 512×512 的設(shè)置上估計(jì) FLOPs 和延遲。由于輸入圖像的分辨率較高，在表9和表10中，GPU 延遲在測(cè)量時(shí)使用了大小為2的 Batch Size。在圖8中，作者將 FastViT 與最近的工作進(jìn)行了比較。FastViT-MA36 的 mIoU 比 PoolFormer-M36 高 5.2%，但是 PoolFormer 具有更高的 FLOPs、參數(shù)量和延遲。

圖8：語(yǔ)義分割實(shí)驗(yàn)結(jié)果

目標(biāo)檢測(cè)和實(shí)例分割實(shí)驗(yàn)實(shí)驗(yàn) MS-COCO 數(shù)據(jù)集，實(shí)驗(yàn)結(jié)果如下圖9所示。所有模型都使用 Mask-RCNN 目標(biāo)檢測(cè)和實(shí)例分割頭按照 1x schedule 進(jìn)行訓(xùn)練。所有的模型都是用預(yù)先訓(xùn)練好的對(duì)應(yīng)圖像分類(lèi)模型的權(quán)重進(jìn)行初始化。結(jié)果顯示出 FastViT 在多種延遲機(jī)制下實(shí)現(xiàn)了最先進(jìn)的性能。FastViT-MA36 模型的性能與 CMT-S 相似，但在桌面GPU 和移動(dòng)設(shè)備上分別快2.4倍和4.3倍。

圖9：目標(biāo)檢測(cè)和實(shí)例分割實(shí)驗(yàn)結(jié)果

總結(jié)

本文提出了一種通用的 CNN 和 Transformer 混合的視覺(jué)基礎(chǔ)模型，是由 MobileOne 原班人馬打造，可以看做是 MobileOne 的方法在 Transformer 上的一個(gè)改進(jìn)型的應(yīng)用。作者引入了一種新的 token mixer，叫做 RepMixer，它使用結(jié)構(gòu)重新參數(shù)化技術(shù)，通過(guò)刪除網(wǎng)絡(luò)中的 Shortcut 來(lái)降低內(nèi)存訪(fǎng)問(wèn)成本，尤其是在較高分辨率時(shí)。作者還提出了進(jìn)一步的架構(gòu)更改，以提高 ImageNet 分類(lèi)任務(wù)和其他下游任務(wù)的性能。在移動(dòng)設(shè)備和 ImageNet 數(shù)據(jù)集上的精度相同的前提下，F(xiàn)astViT 比 CMT 快3.5倍，比 EfficientNet 快4.9倍，比 ConvNeXt 快1.9倍。在類(lèi)似的延遲下，F(xiàn)astViT 在 ImageNet 上獲得的 Top-1 準(zhǔn)確率比 MobileOne 高 4.2%，是一種極具競(jìng)爭(zhēng)力的混合架構(gòu)模型。