論文地址：https://arxiv.org/abs/2307.09283

代碼地址：https://github.com/THU-MIG/RepViT

可以看出，RepViT 相比于其它主流的移動(dòng)端 ViT 架構(gòu)確實(shí)時(shí)很優(yōu)異。接下來讓我們來看下本工作做了哪些貢獻(xiàn)：

1. 文中提到，輕量級(jí) ViTs 通常比輕量級(jí) CNNs 在視覺任務(wù)上表現(xiàn)得更好，這主要?dú)w功于它們的多頭自注意力模塊(MSHA)可以讓模型學(xué)習(xí)全局表示。然而，輕量級(jí) ViTs 和輕量級(jí) CNNs 之間的架構(gòu)差異尚未得到充分研究。
2. 在這項(xiàng)研究中，作者們通過整合輕量級(jí) ViTs 的有效架構(gòu)選擇，逐步提升了標(biāo)準(zhǔn)輕量級(jí) CNN（特別是 MobileNetV3 的移動(dòng)友好性。這便衍生出一個(gè)新的純輕量級(jí) CNN 家族的誕生，即RepViT。值得注意的是，盡管 RepViT 具有 MetaFormer 結(jié)構(gòu)，但它完全由卷積組成。
3. 實(shí)驗(yàn)結(jié)果表明，RepViT 超越了現(xiàn)有的最先進(jìn)的輕量級(jí) ViTs，并在各種視覺任務(wù)上顯示出優(yōu)于現(xiàn)有最先進(jìn)輕量級(jí)ViTs的性能和效率，包括 ImageNet 分類、COCO-2017 上的目標(biāo)檢測(cè)和實(shí)例分割，以及 ADE20k 上的語義分割。特別地，在ImageNet上，RepViT 在 iPhone 12 上達(dá)到了近乎 1ms 的延遲和超過 80% 的Top-1 準(zhǔn)確率，這是輕量級(jí)模型的首次突破。

好了，接下來大家應(yīng)該關(guān)心的應(yīng)該時(shí)“如何設(shè)計(jì)到如此低延遲但精度還很6的模型”出來呢？

方法

再 ConvNeXt 中，作者們是基于 ResNet50 架構(gòu)的基礎(chǔ)上通過嚴(yán)謹(jǐn)?shù)睦碚摵蛯?shí)驗(yàn)分析，最終設(shè)計(jì)出一個(gè)非常優(yōu)異的足以媲美 Swin-Transformer 的純卷積神經(jīng)網(wǎng)絡(luò)架構(gòu)。同樣地，RepViT也是主要通過將輕量級(jí) ViTs 的架構(gòu)設(shè)計(jì)逐步整合到標(biāo)準(zhǔn)輕量級(jí) CNN，即MobileNetV3-L，來對(duì)其進(jìn)行針對(duì)性地改造（魔改）。在這個(gè)過程中，作者們考慮了不同粒度級(jí)別的設(shè)計(jì)元素，并通過一系列步驟達(dá)到優(yōu)化的目標(biāo)。

訓(xùn)練配方的對(duì)齊

首先，論文中引入了一種衡量移動(dòng)設(shè)備上延遲的指標(biāo)，并將訓(xùn)練策略與現(xiàn)有的輕量級(jí) ViTs 對(duì)齊。這一步驟主要是為了確保模型訓(xùn)練的一致性，其涉及兩個(gè)概念，即延遲度量和訓(xùn)練策略的調(diào)整。

延遲度量指標(biāo)

為了更準(zhǔn)確地衡量模型在真實(shí)移動(dòng)設(shè)備上的性能，作者選擇了直接測(cè)量模型在設(shè)備上的實(shí)際延遲，以此作為基準(zhǔn)度量。這個(gè)度量方法不同于之前的研究，它們主要通過FLOPs或模型大小等指標(biāo)優(yōu)化模型的推理速度，這些指標(biāo)并不總能很好地反映在移動(dòng)應(yīng)用中的實(shí)際延遲。

訓(xùn)練策略的對(duì)齊

這里，將 MobileNetV3-L 的訓(xùn)練策略調(diào)整以與其他輕量級(jí) ViTs 模型對(duì)齊。這包括使用 AdamW 優(yōu)化器【ViTs 模型必備的優(yōu)化器】，進(jìn)行 5 個(gè) epoch 的預(yù)熱訓(xùn)練，以及使用余弦退火學(xué)習(xí)率調(diào)度進(jìn)行 300 個(gè) epoch 的訓(xùn)練。盡管這種調(diào)整導(dǎo)致了模型準(zhǔn)確率的略微下降，但可以保證公平性。

塊設(shè)計(jì)的優(yōu)化

接下來，基于一致的訓(xùn)練設(shè)置，作者們探索了最優(yōu)的塊設(shè)計(jì)。塊設(shè)計(jì)是 CNN 架構(gòu)中的一個(gè)重要組成部分，優(yōu)化塊設(shè)計(jì)有助于提高網(wǎng)絡(luò)的性能。

分離 Token 混合器和通道混合器

這塊主要是對(duì) MobileNetV3-L 的塊結(jié)構(gòu)進(jìn)行了改進(jìn)，分離了令牌混合器和通道混合器。原來的 MobileNetV3 塊結(jié)構(gòu)包含一個(gè) 1x1 擴(kuò)張卷積，然后是一個(gè)深度卷積和一個(gè) 1x1 的投影層，然后通過殘差連接連接輸入和輸出。在此基礎(chǔ)上，RepViT 將深度卷積提前，使得通道混合器和令牌混合器能夠被分開。為了提高性能，還引入了結(jié)構(gòu)重參數(shù)化來在訓(xùn)練時(shí)為深度濾波器引入多分支拓?fù)洹Ｗ罱K，作者們成功地在 MobileNetV3 塊中分離了令牌混合器和通道混合器，并將這種塊命名為 RepViT 塊。

降低擴(kuò)張比例并增加寬度

在通道混合器中，原本的擴(kuò)張比例是 4，這意味著 MLP 塊的隱藏維度是輸入維度的四倍，消耗了大量的計(jì)算資源，對(duì)推理時(shí)間有很大的影響。為了緩解這個(gè)問題，我們可以將擴(kuò)張比例降低到 2，從而減少了參數(shù)冗余和延遲，使得 MobileNetV3-L 的延遲降低到 0.65ms。隨后，通過增加網(wǎng)絡(luò)的寬度，即增加各階段的通道數(shù)量，Top-1 準(zhǔn)確率提高到 73.5%，而延遲只增加到 0.89ms！

宏觀架構(gòu)元素的優(yōu)化

在這一步，本文進(jìn)一步優(yōu)化了MobileNetV3-L在移動(dòng)設(shè)備上的性能，主要是從宏觀架構(gòu)元素出發(fā)，包括 stem，降采樣層，分類器以及整體階段比例。通過優(yōu)化這些宏觀架構(gòu)元素，模型的性能可以得到顯著提高。

淺層網(wǎng)絡(luò)使用卷積提取器

圖片

ViTs 通常使用一個(gè)將輸入圖像分割成非重疊補(bǔ)丁的 "patchify" 操作作為 stem。然而，這種方法在訓(xùn)練優(yōu)化性和對(duì)訓(xùn)練配方的敏感性上存在問題。因此，作者們采用了早期卷積來代替，這種方法已經(jīng)被許多輕量級(jí) ViTs 所采納。對(duì)比之下，MobileNetV3-L 使用了一個(gè)更復(fù)雜的 stem 進(jìn)行 4x 下采樣。這樣一來，雖然濾波器的初始數(shù)量增加到24，但總的延遲降低到0.86ms，同時(shí) top-1 準(zhǔn)確率提高到 73.9%。

更深的下采樣層

在 ViTs 中，空間下采樣通常通過一個(gè)單獨(dú)的補(bǔ)丁合并層來實(shí)現(xiàn)。因此這里我們可以采用一個(gè)單獨(dú)和更深的下采樣層，以增加網(wǎng)絡(luò)深度并減少由于分辨率降低帶來的信息損失。具體地，作者們首先使用一個(gè) 1x1 卷積來調(diào)整通道維度，然后將兩個(gè) 1x1 卷積的輸入和輸出通過殘差連接，形成一個(gè)前饋網(wǎng)絡(luò)。此外，他們還在前面增加了一個(gè) RepViT 塊以進(jìn)一步加深下采樣層，這一步提高了 top-1 準(zhǔn)確率到 75.4%，同時(shí)延遲為 0.96ms。

更簡(jiǎn)單的分類器

在輕量級(jí) ViTs 中，分類器通常由一個(gè)全局平均池化層后跟一個(gè)線性層組成。相比之下，MobileNetV3-L 使用了一個(gè)更復(fù)雜的分類器。因?yàn)楝F(xiàn)在最后的階段有更多的通道，所以作者們將它替換為一個(gè)簡(jiǎn)單的分類器，即一個(gè)全局平均池化層和一個(gè)線性層，這一步將延遲降低到 0.77ms，同時(shí) top-1 準(zhǔn)確率為 74.8%。

整體階段比例

階段比例代表了不同階段中塊數(shù)量的比例，從而表示了計(jì)算在各階段中的分布。論文選擇了一個(gè)更優(yōu)的階段比例 1:1:7:1，然后增加網(wǎng)絡(luò)深度到 2:2:14:2，從而實(shí)現(xiàn)了一個(gè)更深的布局。這一步將 top-1 準(zhǔn)確率提高到 76.9%，同時(shí)延遲為 1.02 ms。

微觀設(shè)計(jì)的調(diào)整

接下來，RepViT 通過逐層微觀設(shè)計(jì)來調(diào)整輕量級(jí) CNN，這包括選擇合適的卷積核大小和優(yōu)化擠壓-激勵(lì)（Squeeze-and-excitation，簡(jiǎn)稱SE）層的位置。這兩種方法都能顯著改善模型性能。

卷積核大小的選擇

眾所周知，CNNs 的性能和延遲通常受到卷積核大小的影響。例如，為了建模像 MHSA 這樣的遠(yuǎn)距離上下文依賴，ConvNeXt 使用了大卷積核，從而實(shí)現(xiàn)了顯著的性能提升。然而，大卷積核對(duì)于移動(dòng)設(shè)備并不友好，因?yàn)樗挠?jì)算復(fù)雜性和內(nèi)存訪問成本。MobileNetV3-L 主要使用 3x3 的卷積，有一部分塊中使用 5x5 的卷積。作者們將它們替換為3x3的卷積，這導(dǎo)致延遲降低到 1.00ms，同時(shí)保持了76.9%的top-1準(zhǔn)確率。

SE 層的位置

自注意力模塊相對(duì)于卷積的一個(gè)優(yōu)點(diǎn)是根據(jù)輸入調(diào)整權(quán)重的能力，這被稱為數(shù)據(jù)驅(qū)動(dòng)屬性。作為一個(gè)通道注意力模塊，SE層可以彌補(bǔ)卷積在缺乏數(shù)據(jù)驅(qū)動(dòng)屬性上的限制，從而帶來更好的性能。MobileNetV3-L 在某些塊中加入了SE層，主要集中在后兩個(gè)階段。然而，與分辨率較高的階段相比，分辨率較低的階段從SE提供的全局平均池化操作中獲得的準(zhǔn)確率提升較小。作者們?cè)O(shè)計(jì)了一種策略，在所有階段以交叉塊的方式使用SE層，從而在最小的延遲增量下最大化準(zhǔn)確率的提升，這一步將top-1準(zhǔn)確率提升到77.4%，同時(shí)延遲降低到0.87ms。【這一點(diǎn)其實(shí)百度在很早前就已經(jīng)做過實(shí)驗(yàn)比對(duì)得到過這個(gè)結(jié)論了，SE 層放置在靠近深層的地方效果好】

網(wǎng)絡(luò)架構(gòu)

最終，通過整合上述改進(jìn)策略，我們便得到了模型RepViT的整體架構(gòu)，該模型有多個(gè)變種，例如RepViT-M1/M2/M3。同樣地，不同的變種主要通過每個(gè)階段的通道數(shù)和塊數(shù)來區(qū)分。

實(shí)驗(yàn)

圖像分類

檢測(cè)與分割

總結(jié)

本文通過引入輕量級(jí) ViT 的架構(gòu)選擇，重新審視了輕量級(jí) CNNs 的高效設(shè)計(jì)。這導(dǎo)致了 RepViT 的出現(xiàn)，這是一種新的輕量級(jí) CNNs 家族，專為資源受限的移動(dòng)設(shè)備設(shè)計(jì)。在各種視覺任務(wù)上，RepViT 超越了現(xiàn)有的最先進(jìn)的輕量級(jí) ViTs 和 CNNs，顯示出優(yōu)越的性能和延遲。這突顯了純粹的輕量級(jí) CNNs 對(duì)移動(dòng)設(shè)備的潛力。