Transformer 目前在 CV 領(lǐng)域愈發(fā)火熱，這份火熱促使著優(yōu)秀學(xué)者們思考一個更深層次的問題。部分學(xué)者認為 Transformer 之所以 work 更加本質(zhì)的原因在于其大的感受野（論文直達）。根據(jù)有效感受野（ERF）理論，ERF 大小與 kernel 大小成正比關(guān)系，與模型深度的平方根也成正比關(guān)系。 所以通過堆疊層數(shù)實現(xiàn)大感受野必然不如增加卷積 kernel 大小更高效。因此有學(xué)者提出超大 kernel 卷積的網(wǎng)絡(luò)結(jié)構(gòu)，并證明在目標檢測和語義分割等任務(wù)上超過 Swin Transformer 而且遠超傳統(tǒng)小卷積模型。

什么是大 kernel，什么是 depthwise 卷積？

CNN 中最常見的卷積 kernel 大小有 2x2, 3x3, 5x5, 7x7 等，在本文中我們將卷積 kernel 大小超過 9x9 的視作大 kernel，同時以下所有數(shù)據(jù)都是近似數(shù)據(jù)。我們不難看出隨著卷積 kernel 大小的增加，卷積的參數(shù)量和計算量都呈平方增長，這往往也是大家不喜歡用大 kernel 卷積的其中一個原因。為了獲得大 kernel 卷積帶來的收益的同時降低其計算量和參數(shù)量，我們一般將大 kernel 卷積設(shè)計成 depthwise 卷積。如下圖所示，depthwise 卷積通過逐通道（channel） 做卷積，可以將計算量和參數(shù)量降低到 Dense 卷積的 input channel 分之一。

大 kernel depthwise 卷積為什么值得優(yōu)化？

Roofline Model

為了解釋清楚為什么大 kernel 值得優(yōu)化這個問題，我們需要借助 Roofline 模型的幫助。如下圖所示，Roofline 嘗試解釋一件非常簡單的事情，即應(yīng)用在特定計算設(shè)備下能達到多快的計算速度。

• 理論峰值 TP：描述了計算設(shè)備的性能上限，指的是一個計算設(shè)備每秒鐘最多所能完成的浮點運算數(shù)，單位是?FLOPS。

• *最大帶寬 B8：描述計算設(shè)備的帶寬上限，指的是一個計算設(shè)備每秒最多所能完成的內(nèi)存交換量，單位是?Byte/s。

• 最大計算密度 IM：描述計算設(shè)備單位內(nèi)存交換最多用來進行多少次運算，單位是?FLOPs/Byte。

"Roofline" 指的是由計算設(shè)備理論算力峰值和最大訪存帶寬這兩個參數(shù)所決定的“屋頂”形態(tài)。其中設(shè)備理論峰值決定“屋頂”的高度（藍色線段），設(shè)備最大訪存帶寬決定了“屋檐”的斜率（紅色線段）。Roofline 模型劃分出來兩個瓶頸區(qū)域，分別為 Compute Bound 和 Memory Bound。

當應(yīng)用的計算密度 I 超過最大計算密度 IM 時，此時無論應(yīng)用的計算密度多大，它的性能最高只能達到計算設(shè)備的理論峰值 TP。此時應(yīng)用的性能 P 被設(shè)備理論峰值限制無法和計算密度 I 成正比，所以叫做 Compute Bound。當應(yīng)用的計算密度 I 小于最大計算密度 IM 時，此時性能 P 將由設(shè)備最大帶寬和應(yīng)用計算密度決定。不難看出對于處在 Memory Bound 區(qū)間的應(yīng)用，增加設(shè)備帶寬和增加計算密度可以使應(yīng)用性能達到線性增長的目的。

走出對 depthwise 卷積速度的 "思維誤區(qū)"

為什么不是大 kernel Dense 卷積

現(xiàn)如今針對 Dense 卷積我們已經(jīng)有了包括 Direct、im2col/implicit GEMM、Winograd 和 FFT 等多種優(yōu)化手段，可以說已經(jīng)足夠成熟了?？墒侨绻覀儝侀_模型參數(shù)量，僅僅從運行效率的角度思考一個問題，為什么我們不用大 kernel Dense 卷積而選擇大 kernel depthwise 卷積呢？

為了探尋這個問題的答案，我們結(jié)合 Roofline 模型具體分析。本文選取 2080Ti 顯卡為計算設(shè)備，它的實測 L2 cache 帶寬為2.16TB/s，理論峰值性能為 4352 FFMA Cores * 1.545 GHZ * 2 = 13.447 TFLOPS。我們假設(shè) CUDA 中每個 thread 負責(zé)計算的 output 數(shù)據(jù)都放在寄存器中累加，我們假設(shè) L1 cache 100% 命中，忽略寫回 output 的過程。由于現(xiàn)代計算設(shè)備的設(shè)計足夠合理，實際卷積計算中足以抵消很多耗時較長的訪存操作，同時為了簡化分析復(fù)雜度，在這里我們假設(shè) L2 cache 100% 命中，使用 L2 cache 的最大帶寬作為分析參數(shù)。本文使用的卷積輸入 shape 是（n, ic, ih, iw），kernel 是 （oc, ic, kh, kw），output 是 （n, oc, oh, ow）。

對 Dense 卷積而言，一種通用優(yōu)化計算手段就是 im2col/implicit GEMM。由于其太經(jīng)典了我們在這里不再贅述 im2col 的過程，感興趣的可以翻閱我們之前寫的文章《MegEngine TensorCore 卷積算子實現(xiàn)原理》。在經(jīng)過了 im2col 變換之后，我們就成功的將卷積轉(zhuǎn)換成了矩陣乘的形式。其中矩陣乘的 M = oc, N = n*oh*ow, K = ic*kh*kw，具體如下圖所示。

對于矩陣乘特別是大規(guī)模矩陣乘，cuBlas 等計算庫已經(jīng)優(yōu)化的足夠好了，基本上可以接近設(shè)備理論峰值，這里我們結(jié)合 Roofline 簡單分析一下性能。為了充分適應(yīng)硬件體系結(jié)構(gòu)特征，充分利用多級存儲增大訪存帶寬，我們需要對矩陣乘進行分塊計算。如下圖所示，假如 cuda 中每個 Thread Block 處理 BMxBN 的 output，此時 kernel 分塊大小為 BMxBK，input 分塊大小為 BKxBN。則計算量為 BM*BN*BK*2，訪存量為 (BM*BK + BN*BK)*4。計算密度為?。按照 Roofline 模型的描述，計算設(shè)備的

FLOPs/Byte，若要達到設(shè)備理論峰值我們只要保證計算密度大于 IM 即可。如果我們按照 BM=32, BN=32 來算的話，則此時的計算密度將達到 8 FLOPs/Byte，顯然是大于 IM 的。此時如果忽略 TP 的限制假如打滿設(shè)備最大帶寬，最大可能達到的性能 P = 8*2.16 = 17.28 TFLOPS。結(jié)合 Roofline 模型不難看出此時處于 Compute Bound 區(qū)域。由于 Compute Bound 區(qū)域的計算速度已經(jīng)接近理論峰值，已經(jīng)不能增加了。如果我們采用大 kernel 的話，隨著 kernel size 的增加計算量會呈平方增長，所以相應(yīng)的運行時間也會隨之增長，這顯然是不可接受的。

depthwise 卷積速度的“騙局”

對 Dense 卷積分析讓我們得到了一個結(jié)論即 “隨著 kernel 的增大，卷積時間呈平方增長”。很多人想當然的將這個結(jié)論平移到了 depthwise 卷積上，這其實是一種思維誤區(qū)。

讓我們同樣嘗試用 im2col/implicit GEMM 的方法分析 depthwise 卷積。由于 depthwise 是逐 channel 做卷積的，所以可以看做 channel 數(shù)量的單通道卷積。在經(jīng)過 im2col 變換之后我們將獲得一個 Batched GEMV，每個 batch 的 GEMV 如下圖所示。

如果我們保持和 Dense 卷積一樣的分塊策略的話，每個 batch 的 GEMV 如下圖所示。相應(yīng)的此時的計算密度為?。先不說這是一個 Batched GEMV，單獨看一個 GEMV 也不難發(fā)現(xiàn)此時的計算密度是很差的，BN = 1 時最高大概能達到 0.25 FLOPs/Byte，相應(yīng)的最大達到的性能 P = 0.25*2.16 = 0.54 TFLOPS。當然了實際應(yīng)用中 GEMV 還有其他計算方式，我們的分析方法就不一定準確了。但此處想表達的意思是 Batched GEMV 比 GEMM 更難優(yōu)化。假如 kernel 為 3x3，此時 M=1, K=9, N 受限于 oh 和 ow 也不會很大，此時的 GEMV 性能肯定遠達不到峰值，并且 GEMV 也不能利用 TensorCore 加速。

如果我們嘗試使用 Direct 的方式處理 depthwise 卷積的話會不會好一點呢？例如我們讓 cuda 中每個 warp 32 個線程負責(zé)計算 ohxow 的輸出，kernel size 為 khxkw，此時：

• 計算量 = oh*ow*kh*kw*2 FLOPs

• 訪存量 = (kh*kw + (oh+kh-1)*(ow+kw-1)) * 4 Bytes，分別為

• kernel: kh*kw

• input: (oh+kh-1)*(ow+kw-1)

• 計算密度為

我們以一個更具體的例子分析，假如我們讓每個 thread 負責(zé)計算 4 個 output 的話，則一個 warp 負責(zé)計算 4x32 的 output，以 kernel （3, 3) 為例。則計算密度為??FLOPs/Byte，最大可達到的性能為 2.16*2.7 = 5.84 TFLOPS，相比于理論峰值 13.447 TFLOPS 仍有很大差距。雖然增加 output 能繼續(xù)增加計算密度，但是受限于卷積本身的輸出大小和每個 SM 中有限的 register file 等計算資源，每個 warp 計算的 output 并不能無限增加。這也是 depthwise 卷積需要更加仔細的優(yōu)化，否則一不小心性能就會很差的其中一個原因。

綜合 im2col 和 Direct 兩個方面的分析結(jié)論，我們認識到和 Dense 卷積不同的是 depthwise 卷積很多時候是一個 Memory Bound 的操作。而結(jié)合 Roofline 模型對 Memory Bound 瓶頸的分析和建議，此時增加計算密度和增加帶寬都可以增加性能。在固定設(shè)備的情況下我們無法增加帶寬了，所以看起來增加計算密度是一個可行的方案。通過觀察計算密度公式我們不難發(fā)現(xiàn)，增加 depthwise 卷積的 kernel size 就是一個增加其計算密度的有效方案，例如保持每個 warp 4x32 的輸出配置下 kernel size 31x31 的 depthwise 卷積計算密度將達到FLOPs/Byte，不難看出此時已經(jīng)變成了 Compute Bound 的操作。

綜上所述，增加卷積 kernel size 會使得計算量增加。同時因為 Dense 卷積處于 Compute Bound 區(qū)域，所以其運行速度受限于設(shè)備理論峰值無法提升，因此針對 Dense 卷積我們不難歸納出** “隨著 kernel 的增大，卷積時間呈平方增長”** 的規(guī)律。但是 depthwise 卷積是一種 Memory Bound 的操作，而隨著 kernel size 的增加其計算密度也會增大，所以其運行性能也會隨之增大。此時的卷積的運行時間并不會顯著增長，所以它并不適用?“隨著 kernel 的增大，卷積時間呈平方增長”?這個結(jié)論。這也是我們認為大 kernel depthwise 還有較大的優(yōu)化潛力，其運行時間并不會明顯差于小 kernel depthwise 卷積的依據(jù)。

現(xiàn)有優(yōu)化方法為什么不行？

上一節(jié)我們已經(jīng)解釋了為什么 im2col/implicit GEMM 不適合 depthwise 卷積，direct 也需要付出很大精力才能寫好。另外，提到大 kernel 則不能不提?FFT 算法，但 FFT 在計算 depthwise 卷積的時候只能逐通道計算，性能不如預(yù)期。并且 FFT 有其缺陷例如精度問題，對半精度計算并不友好，也不能被量化。我們在 2080Ti 上使用 input 和 output 形狀都是 (n, c, h, w) = (64, 384, 32, 32) 的用例對 cudnn做了一次測速，我們遍歷所有的 cudnn 算子（內(nèi)含 FFT）并選擇最快的那個算子進行測試。結(jié)果如下：

在大 kernel size 下 cudnn 的表現(xiàn)很差，主要原因是 cudnn 沒有針對性優(yōu)化。我們注意到很多時候 cudnn 調(diào)用到了內(nèi)部的 implicit_gemm 實現(xiàn)，這不利于發(fā)揮設(shè)備的計算性能。因為對于 depthwise 卷積而言，im2col 之后將會是一個 batch = channel，M = 1，N=nhw， K = kh*kw 的 batched GEMV，這種情況也很難打滿設(shè)備峰值。

MegEngine 的優(yōu)化效果和簡單分析

鑒于以上分析，大 kernel depthwise 卷積有很大的優(yōu)化潛力，所以 MegEngine 緊跟學(xué)界動態(tài)對大 kernel depthwise 卷積進行了深度優(yōu)化。如上圖所示，經(jīng)過我們的優(yōu)化后，隨著 kernel size 的增加，算子性能基本呈現(xiàn)線性增長的趨勢，部分情況下算子可以逼近硬件的單精度浮點理論峰值。

如下圖所示，優(yōu)化后的大 kernel depthwise 卷積比 PyTorch 快 10.x 倍，代碼附在文末，感興趣的同學(xué)歡迎來體驗一把。而且我們不難發(fā)現(xiàn)，隨著 kernel size 的增加模型訓(xùn)練時間并沒有顯著增加。原因就在于 kernel size 不夠大的時候算子處于 Memory Bound 狀態(tài)，遠沒有達到理論峰值，此時增加計算密度反而不會對算子運行時間造成很大影響。

想知道 MegEngine 是如何將 31*31 的 DWconv 優(yōu)化快了 10 余倍？還有 ConvNext，RepLKNet 為何不約而同將 kernel size 增大，更大的 kernel size 到底給模型帶來了什么？來 MegEngine Meetup 一起聊聊吧。

Meetup 視頻回顧

北京時間 3.19 上午 10:00，MegEngine Meetup 圍繞“Large Kernel Makes CNN Great Again”主題，帶來了精彩線上分享，完整視頻回顧見：https://www.bilibili.com/video/BV16Y411n7bP

附：測試代碼

MegEngine 測試代碼

PyTorch 測試代碼

paper 地址：https://arxiv.org/abs/2203.06717

MegEngin 復(fù)現(xiàn)代碼:https://github.com/megvii-research/RepLKNet

GitHub：MegEngine 天元

官網(wǎng)：MegEngine-深度學(xué)習(xí)，簡單開發(fā)

歡迎加入 MegEngine 技術(shù)交流 QQ 群：1029741705