如何擁有一個兼具CNN的速度、Transformer精度的模型？字節(jié)甩出TRT-ViT教你

論文鏈接：https://arxiv.org/abs/2205.09579

本文作者從實(shí)際應(yīng)用的角度重新審視現(xiàn)有的Transformer。它們中的大多數(shù)甚至不如基本的ResNets系列高效，并且偏離了現(xiàn)實(shí)的部署場景。這可能是由于當(dāng)前衡量計(jì)算效率的標(biāo)準(zhǔn)，例如FLOP或參數(shù)是片面的、次優(yōu)的同時對硬件也不敏感的。

因此，本文直接將特定硬件上的TensorRT延遲作為效率指標(biāo)，提供了更全面的計(jì)算能力、內(nèi)存成本和帶寬反饋。在一系列對照實(shí)驗(yàn)的基礎(chǔ)上，本文得出了TensorRT的4個設(shè)計(jì)指南，例如Stage-level的Early CNN和Late Transformer，Block-level的Early Transformer和Late CNN。

據(jù)此，作者提出了一個面向TensorRT的Transformer家族，簡稱TRT-ViT。大量實(shí)驗(yàn)表明，TRT-ViT在各種視覺任務(wù)（例如圖像分類、目標(biāo)檢測和語義分割）的延遲/準(zhǔn)確性權(quán)衡方面明顯優(yōu)于現(xiàn)有的 ConvNet和視覺Transformer。

例如，在ImageNet-1k top-1 準(zhǔn)確率為82.7% 時，TRT-ViT比CSWin快2.7倍，比Twins快2.0倍。在MS-COCO目標(biāo)檢測任務(wù)上，TRT-ViT實(shí)現(xiàn)了與Twins相當(dāng)?shù)男阅埽瑫r推理速度提高了2.8倍。

簡介

最近，Vision Transformer在圖像分類、語義分割和目標(biāo)檢測等各種計(jì)算機(jī)視覺任務(wù)中取得了顯著的成功，并取得了明顯優(yōu)于CNN的性能提升。然而，從真實(shí)場景的角度來看，CNN仍然主導(dǎo)著視覺架構(gòu)部署。

當(dāng)深入挖掘背后的實(shí)質(zhì)性性能改進(jìn)時現(xiàn)有的Transformer，觀察到這些收益是以大量資源開銷為代價的。此外，當(dāng)Transformer部署在資源受限的設(shè)備上時，它已成為一個障礙。為了減輕Transformer笨重的計(jì)算成本，一系列工作重點(diǎn)開發(fā)了高效的Vision Transformer，如Swin Transformer、PVT、Twins、CoAtNet和MobileViT。

盡管上述工作從不同的角度展示了ViT的效率，但關(guān)鍵設(shè)計(jì)主要由計(jì)算效率的間接標(biāo)準(zhǔn)指導(dǎo)，例如FLOP或參數(shù)。這種迂回的標(biāo)準(zhǔn)是片面的、次優(yōu)的，并且偏離了真實(shí)場景的部署。事實(shí)上，該模型必須應(yīng)對部署過程中的環(huán)境不確定性，這涉及到硬件感知特性，例如內(nèi)存訪問成本和I/O吞吐量。

具體來說，正如許多論文中所述，計(jì)算復(fù)雜度的間接度量在神經(jīng)網(wǎng)絡(luò)架構(gòu)設(shè)計(jì)中被廣泛使用，F(xiàn)LOPs和參數(shù)是近似的，但通常不能真實(shí)反映直接度量。真正關(guān)心的是速度或延遲。

在本文中，為了解決這種差異，作者提供了一種觀點(diǎn)，將特定硬件上的TensorRT延遲視為直接的效率反饋。TensorRT在實(shí)際場景中已成為通用且易于部署的解決方案，有助于提供令人信服的面向硬件的解決方案指導(dǎo)。在這種直接準(zhǔn)確的指導(dǎo)下，重新繪制了圖1中幾種現(xiàn)有競爭模型的準(zhǔn)確性和效率權(quán)衡圖。

如圖1所示，Transformer具有性能良好的優(yōu)勢，而CNN則成功地實(shí)現(xiàn)了高效率。盡管ResNet在性能上不如競爭對手的Transformers令人印象深刻，但它仍然是準(zhǔn)確率-延遲權(quán)衡下的最佳架構(gòu)。

例如，ResNet使用 11.7毫秒（batch size = 8）來達(dá)到81.7%的準(zhǔn)確率。雖然twins-pcpvt-s 在Imagenet-1k分類中達(dá)到了令人印象深刻的準(zhǔn)確率，達(dá)到83.4%，但實(shí)現(xiàn)這一目標(biāo)需要39.8毫秒。這些觀察促使提出一個問題：如何設(shè)計(jì)一個性能與Transformer一樣好、預(yù)測與ResNet一樣快的模型？

為了回答這個問題，作者系統(tǒng)地探索了CNN和Transformer的混合設(shè)計(jì)。遵循Stage到Block的分層路線圖來研究面向TensorRT的架構(gòu)。通過廣泛的實(shí)驗(yàn)，總結(jié)了在TensorRT上設(shè)計(jì)高效網(wǎng)絡(luò)的4個設(shè)計(jì)準(zhǔn)則：

1. Stage-Level：在后期Stage使用Transformer Block可以最大化效率和性能；
2. Stage-Level：由淺入深的Stage模式可以提升性能；
3. Block-Level：Transformer和BottleNeck混合Block比Transformer效率更高；
4. Block-Level：先全局后局部的模式有助于彌補(bǔ)性能問題。

基于上述準(zhǔn)則設(shè)計(jì)了一系列面向TensorRT的Transformer（縮寫為TRT-ViT），由具有ConvNets和Transformers的混合網(wǎng)絡(luò)組成。CNN和Transformer在TRT-ViT的Stage和Block位置互補(bǔ)。

此外，提出了各種TRT-ViT Block以將CNN和Transformer組合成一個串行方案，以解耦異構(gòu)概念并提高信息流的效率。所提出的TRT-ViT在延遲準(zhǔn)確度交易方面優(yōu)于現(xiàn)有的ConvNets和Vision Transformer, 例如，在ImageNet-1k top-1 準(zhǔn)確率為 82.7% 時，TRT-ViT比CSWin快2.7倍，比Twins快2.0倍。

更重要的是，TRT-ViT在下游任務(wù)上表現(xiàn)出更顯著的延遲/準(zhǔn)確性權(quán)衡增益。如圖1所示，TRT-ViT 實(shí)現(xiàn)了與Twins相當(dāng)?shù)男阅埽评硭俣仍贛S-COCO目標(biāo)檢測上提高了2.8倍，在ADE20K語義分割任務(wù)上提高了4.4倍。在相似的TensorRT延遲下，TRT-ViT在COCO檢測上的性能比ResNet高3.2AP（40.4到43.6），在ADE20K分割上比ResNet高6.6%（38.86%到45.4%）。

TRT-ViT

2.1 基于TensorRT的設(shè)計(jì)指南

為了研究分析了2個SOTA網(wǎng)絡(luò) ResNet 和 ViT 的運(yùn)行性能。它們在 ImageNet1K 分類任務(wù)上都非常準(zhǔn)確，并且在業(yè)界很流行。雖然只分析了這2個網(wǎng)絡(luò)，但它們代表了目前的趨勢。它們的核心是 BottleNeck Block 和 Transformer Block ，它們也是其他最先進(jìn)網(wǎng)絡(luò)的關(guān)鍵組件，例如 ConvNext 和 Swin。

為了更好地說明經(jīng)驗(yàn)結(jié)果，遵循 RepLKNet，使用操作或塊的計(jì)算密度（以每秒 Tera 浮點(diǎn)操作數(shù)，TeraFLOPS 衡量）來量化其在硬件上的效率，描述如下

從圖 1(a) 和表 1 可以看出，BottleNeck Block 更高效，具有更高的 TeraFLOPS，而 Transformer Block更強(qiáng)大，在 ImageNet 上具有更好的精度。

在本文中，目標(biāo)是設(shè)計(jì)一個實(shí)現(xiàn)類似 ResNet 的效率和類似 Transformer 的性能的網(wǎng)絡(luò)。設(shè)計(jì)一個混合了 BottleNeck 和 Transformer Block的網(wǎng)絡(luò)很簡單，例如 MobileViT。然而，這些論文要么專注于優(yōu)化 FLOP，這是衡量效率的間接標(biāo)準(zhǔn)，要么將大量精力放在移動設(shè)備上。

將 Transformer 簡單地應(yīng)用于 CNN 通常會導(dǎo)致性能和速度下降。因此，總結(jié)了在 TensorRT 上有效設(shè)計(jì)高效網(wǎng)絡(luò)的4個實(shí)用準(zhǔn)則。

準(zhǔn)則1：Transformer Block在后期使用可以最大限度地提高效率和性能之間的平衡

人們認(rèn)為Transformer塊比卷積塊更強(qiáng)大但效率更低，因?yàn)門ransformer旨在建立特征之間的全局連接。相比之下，卷積只捕獲局部信息。在這里，提供了配備TensorRT的詳細(xì)分析。如圖1所示，堆疊Transformer塊（PVT-Medium）在ImageNet1K分類任務(wù)上的 Top-1 準(zhǔn)確率為 81.2%，而BottleNeck塊（ResNet50）的相應(yīng)準(zhǔn)確率僅為 79.0%。然而，如表1所示，在效率方面，BottleNeck模塊在各種輸入分辨率下的TeraFLOPS始終優(yōu)于Transformer模塊。

此外，當(dāng)輸入分辨率降低時，Transformer和BottleNeck之間的效率差距會縮小。具體來說，當(dāng)輸入分辨率為56×56時，Transformer的TeraFLOPS 值只有81，不到BottleNeck的TeraFLOPS值378的4分之1。但是，當(dāng)分辨率降低到7×7時，Transformer的TeraFLOPS值幾乎等于BottleNeck，分別為599和670。這一觀察促使在TensorRT上設(shè)計(jì)高效網(wǎng)絡(luò)時將Transformer模塊置于后期階段以平衡性能和效率。CoATNet也遵循類似的規(guī)則并提出混合了Transformer和CNN 模塊的網(wǎng)絡(luò)。

遵循準(zhǔn)則1，原生解決方案是在后期將BottleNeck塊替換為Transformer塊，在表 1 中稱為 MixNetV（MixNet Valilla）。可以看到MixNetV比ViT更快，比ResNet更準(zhǔn)確。

準(zhǔn)則2：從淺到深的Stage模式可以提高性能

作者研究了各種 SOTA 網(wǎng)絡(luò)（例如 RestNet 和 ViT）中 TeraParams 的屬性，并在表 1 中展示了發(fā)現(xiàn)。BottleNeck 的 TeraParams 在后期變得更加突出，這表明在后期堆疊更多的 BottleNeck 塊將與早期相比，模型容量更大。

在 ViT 的 Transformer 塊中也可以找到類似的趨勢。這一觀察啟發(fā)使早期階段更淺，晚期階段更深。假設(shè)這種由淺入深的階段模式將帶來更多參數(shù)，提高模型容量并在不降低效率的情況下獲得更好的性能。

為了驗(yàn)證這一點(diǎn)，對ResNet50進(jìn)行了相應(yīng)的修改，并提供了表2中的經(jīng)驗(yàn)結(jié)果。將第1階段和第2階段的Stage深度分別從3、4減少到2、3，同時增加最后階段的Stage深度從3到5。從表2中可以看到Refined-ResNet50在ImageNet1K上的Top-1準(zhǔn)確率比ResNet50高0.3%，同時在TensorRT上稍快。此外，將這種由淺入深的階段模式應(yīng)用到PvT-Small上，得到Refined-PvT-Small，更好更快。最后，在MixNetV上也可以得出類似的結(jié)論，這并不奇怪。

準(zhǔn)則3：BottleNeck Block和Transformer Block混合使用比Transformer更有效

雖然按照準(zhǔn)則 1 和 2 取得了顯著的改進(jìn)，但Refined-MixNetV在效率上無法達(dá)到ResNet的效率。為此，設(shè)計(jì)一個利用Transformer和BottleNeck兩者的優(yōu)點(diǎn)的混合模塊。

在圖3(a)和(b)中提供了2個混合塊，其中首先使用1×1 Conv層來投影輸入通道的數(shù)量，最后使用Concat層來合并全局和局部特征。引入收縮率R來控制Transformer塊在混合塊中的比例，定義為Transformer塊的輸出通道除以混合塊的輸出通道。如圖3(a)所示，在MixBlockA中，創(chuàng)建了2個分支，一個用于Transformer，另一個用于BottleNeck。

此外，通道拆分層用于進(jìn)一步提高效率，其中Transformer塊的輸出通道為，而對應(yīng)的通道為。在MixBlockB中，Transformer塊和BottleNeck塊是順序堆疊的，Transformer塊和BottleNeck塊的輸出通道相等，即此處R=0.5。從表1可以看出，與Transformer塊相比，2個混合塊的TeraFLOPS都顯著增加，這表明混合塊比Transformer塊更有效。而且，它們對應(yīng)的TeraParams也更好，這表明它們具有實(shí)現(xiàn)更好性能的潛力。

用MixBlockA和MixBlockB替換MixNetV中的Transformer塊，構(gòu)建了MixNetA和MixNetB。如表3所示，其中C、T和M分別表示Convolution、Transformer和MixBlock，可以看到MixNetA和MixNetB都優(yōu)于ResNet50。但是，在準(zhǔn)確性方面，它們都無法擊敗MixNetV。此外，PVT-Medium和2種混合網(wǎng)絡(luò)之間的性能差距也不容忽視。

準(zhǔn)則4：Global-then-Local Block有助于彌補(bǔ)性能問題

正如許多論文中所描述的，感受野通常很大的Transformer從輸入特征中提取全局信息，并在每個條目內(nèi)實(shí)現(xiàn)信息交換。相比之下，卷積的小感受野只挖掘局部信息。首先獲取全局信息并在局部進(jìn)行提煉，而不是先提取局部信息并在全局范圍內(nèi)提煉更有意義。此外，MixBlockB中的收縮率R必須為0.5，失去了進(jìn)一步調(diào)整的靈活性。global-then-local塊模式解決了這個問題，其中R可以是0到1之間的任何值。

按照這個準(zhǔn)則，在MixBlockB中轉(zhuǎn)換Transformer操作和Convolution操作的順序，得到MixBlockC，如圖3(c)所示。MixBlockC的TeraFLOPS分?jǐn)?shù)與MixBlockB相當(dāng)。然后，用MixBlockC代替MixBlockB，得到MixNetC。與直覺一致，如表3所示，MixNetC的性能優(yōu)于MixNetB，而不會犧牲任何效率。此外，MixNetC在準(zhǔn)確性和TRT延遲方面都超過了MixNetV。

此外，MixNetC和PVT-Medium之間的差距得到了彌補(bǔ)，MixNetC實(shí)現(xiàn)了類似ResNet的效率和類似Transformer的性能。

2.2 TRT-ViT: An Efficient Network on TensorRT

按照ResNet的基本配置構(gòu)建了TRT-ViT 架構(gòu)，其中采用特征金字塔結(jié)構(gòu)，特征圖的分辨率隨著網(wǎng)絡(luò)深度而降低，而通道數(shù)增加。

接下來，應(yīng)用提出的設(shè)計(jì)準(zhǔn)則來構(gòu)建TRT-ViT。整個架構(gòu)分為5個階段，后期Stage只使用MixBlockC，而卷積層用于早期Stage。此外，使用了由淺入深的Stage模式，與ResNet中的Stage模式相比，早期Stage更淺，后期Stage更深。

為了與其他SOTA網(wǎng)絡(luò)進(jìn)行公平比較，提出了4種變體，TRT-ViT-A/B/C/D，其配置列于表4。3×3表示卷積核大小=3，C表示一個階段的輸出通道。R、S、K是所提出的MixBlockC中的內(nèi)部參數(shù)，其中S代表Transformer Block中的空間縮減率，K表示MixBlockC中BottleNeck的kernel-size大小，R 表示收縮率。

為簡單起見，設(shè)置R=0.5。TRT-ViT-A/B/C/D的Stage深度分別為2-4-5-4、3-4-7-4、3-4-9-6、4-5-9-5。對于所有變體，在最后階段使用MixBlockC，即stage5。

對于更大的型號，TRT-ViT-C/D，另外在stage4使用2個MixBlcokC。另外，每個MLP層的擴(kuò)展率設(shè)置為3，Transformer中的head dim設(shè)置為32。對于標(biāo)準(zhǔn)化和激活函數(shù)，BottleNeck塊使用BatchNorm和ReLU，Transformer塊使用LayerNorm和GeLU。

實(shí)驗(yàn)

3.1 ImageNet-1K

3.2 語義分割

3.3 目標(biāo)檢測與實(shí)例分割

公眾號后臺回復(fù)“CVPR 2022”獲取論文合集打包下載～