為什麼是大 kernel 卷積？

Transformer 目前在 CV 領域愈發火熱，這份火熱促使著優秀學者們思考一個更深層次的問題。部分學者認為 Transformer 之所以 work 更加本質的原因在於其大的感受野（論文直達）。根據有效感受野（ERF）理論，ERF 大小與 kernel 大小成正比關係，與模型深度的平方根也成正比關係。 所以通過堆疊層數實現大感受野必然不如增加捲積 kernel 大小更高效。因此有學者提出超大 kernel 卷積的網路結構，並證明在目標檢測和語義分割等任務上超過 Swin Transformer 而且遠超傳統小卷積模型。

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

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

大 kernel depthwise 卷積為什麼值得優化？

Roofline Model

為了解釋清楚為什麼大 kernel 值得優化這個問題，我們需要藉助 Roofline 模型的幫助。如下圖所示，Roofline 嘗試解釋一件非常簡單的事情，即應用在特定計算裝置下能達到多快的計算速度。

• 理論峰值 TP：描述了計算裝置的效能上限，指的是一個計算裝置每秒鐘最多所能完成的浮點運算數，單位是 FLOPS。
• **最大頻寬 B*8：描述計算裝置的頻寬上限，指的是一個計算裝置每秒最多所能完成的記憶體交換量，單位是 Byte/s。
• 最大計算密度 IM：描述計算裝置單位記憶體交換最多用來進行多少次運算，單位是 FLOPs/Byte
  。

"Roofline" 指的是由計算裝置理論算力峰值和最大訪存頻寬這兩個引數所決定的“屋頂”形態。其中裝置理論峰值決定“屋頂”的高度（藍色線段），裝置最大訪存頻寬決定了“屋簷”的斜率（紅色線段）。Roofline 模型劃分出來兩個瓶頸區域，分別為 Compute Bound 和 Memory Bound。

當應用的計算密度 I 超過最大計算密度 IM 時，此時無論應用的計算密度多大，它的效能最高只能達到計算裝置的理論峰值 TP。此時應用的效能 P 被裝置理論峰值限制無法和計算密度 I 成正比，所以叫做 Compute Bound。當應用的計算密度 I 小於最大計算密度 IM 時，此時效能 P 將由裝置最大頻寬和應用計算密度決定。不難看出對於處在 Memory Bound 區間的應用，增加裝置頻寬和增加計算密度可以使應用效能達到線性增長的目的。

走出對 depthwise 卷積速度的 "思維誤區"

為什麼不是大 kernel Dense 卷積

現如今針對 Dense 卷積我們已經有了包括 Direct、im2col/implicit GEMM、Winograd 和 FFT 等多種優化手段，可以說已經足夠成熟了。可是如果我們拋開模型引數量，僅僅從執行效率的角度思考一個問題，為什麼我們不用大 kernel Dense 卷積而選擇大 kernel depthwise 卷積呢？

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

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

對於矩陣乘特別是大規模矩陣乘，cuBlas 等計算庫已經優化的足夠好了，基本上可以接近裝置理論峰值，這裡我們結合 Roofline 簡單分析一下效能。為了充分適應硬體體系結構特徵，充分利用多級儲存增大訪存頻寬，我們需要對矩陣乘進行分塊計算。如下圖所示，假如 cuda 中每個 Thread Block 處理 BMxBN 的 output，此時 kernel 分塊大小為 BMxBK，input 分塊大小為 BKxBN。則計算量為 BM*BN*BK*2，訪存量為 (BM*BK + BN*BK)*4。計算密度為 \(\frac{BM*BN*2}{(BM+BN)*4}\)

。按照 Roofline 模型的描述，計算裝置的\(IM = \frac{TP}{B} = \frac{13.447}{2.16} = 6.225\)

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

depthwise 卷積速度的“騙局”

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

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

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

如果我們嘗試使用 Direct 的方式處理 depthwise 卷積的話會不會好一點呢？例如我們讓 cuda 中每個 warp 32 個執行緒負責計算 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)
• 計算密度為\(\frac{oh*ow*kh*kw*2}{(kh*kw+(oh+kh-1)*(ow+kw-1))*4}\)
  我們以一個更具體的例子分析，假如我們讓每個 thread 負責計算 4 個 output 的話，則一個 warp 負責計算 4x32 的 output，以 kernel （3, 3) 為例。則計算密度為 $\frac{432332}{(33+634)4} = 2.7 $ FLOPs/Byte，最大可達到的效能為 2.162.7 = 5.84 TFLOPS，相比於理論峰值 13.447 TFLOPS 仍有很大差距。雖然增加 output 能繼續增加計算密度，但是受限於卷積本身的輸出大小和每個 SM 中有限的 register file 等計算資源，每個 warp 計算的 output 並不能無限增加。這也是 depthwise 卷積需要更加仔細的優化，否則一不小心效能就會很差的其中一個原因。

綜合 im2col 和 Direct 兩個方面的分析結論，我們認識到和 Dense 卷積不同的是 depthwise 卷積很多時候是一個 Memory Bound 的操作。而結合 Roofline 模型對 Memory Bound 瓶頸的分析和建議，此時增加計算密度和增加頻寬都可以增加效能。在固定裝置的情況下我們無法增加帶寬了，所以看起來增加計算密度是一個可行的方案。通過觀察計算密度公式我們不難發現，增加 depthwise 卷積的 kernel size 就是一個增加其計算密度的有效方案，例如保持每個 warp 4x32 的輸出配置下 kernel size 31x31 的 depthwise 卷積計算密度將達到\(\frac{4*32*31*31*2}{(31*31+34*62)*4} = 20\) FLOPs/Byte，不難看出此時已經變成了 Compute Bound 的操作。

綜上所述，增加捲積 kernel size 會使得計算量增加。同時因為 Dense 卷積處於 Compute Bound 區域，所以其執行速度受限於裝置理論峰值無法提升，因此針對 Dense 卷積我們不難歸納出** “隨著 kernel 的增大，卷積時間呈平方增長”** 的規律。但是 depthwise 卷積是一種 Memory Bound 的操作，而隨著 kernel size 的增加其計算密度也會增大，所以其執行效能也會隨之增大。此時的卷積的執行時間並不會顯著增長，所以它並不適用 “隨著 kernel 的增大，卷積時間呈平方增長” 這個結論。這也是我們認為大 kernel depthwise 還有較大的優化潛力，其執行時間並不會明顯差於小 kernel depthwise 卷積的依據。

現有優化方法為什麼不行？

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

在大 kernel size 下 cudnn 的表現很差，主要原因是 cudnn 沒有針對性優化。我們注意到很多時候 cudnn 呼叫到了內部的 implicit_gemm 實現，這不利於發揮裝置的計算效能。因為對於 depthwise 卷積而言，im2col 之後將會是一個 batch = channel，M = 1，N=nhw， K = kh*kw 的 batched GEMV，這種情況也很難打滿裝置峰值。

MegEngine 的優化效果和簡單分析

鑑於以上分析，大 kernel depthwise 卷積有很大的優化潛力，所以 MegEngine 緊跟學界動態對大 kernel depthwise 卷積進行了深度優化。如上圖所示，經過我們的優化後，隨著 kernel size 的增加，運算元效能基本呈現線性增長的趨勢，部分情況下運算元可以逼近硬體的單精度浮點理論峰值。

如下圖所示，優化後的大 kernel depthwise 卷積比 PyTorch 快 10.x 倍，程式碼附在文末，感興趣的同學歡迎來體驗一把。而且我們不難發現，隨著 kernel size 的增加模型訓練時間並沒有顯著增加。原因就在於 kernel size 不夠大的時候運算元處於 Memory Bound 狀態，遠沒有達到理論峰值，此時增加計算密度反而不會對運算元執行時間造成很大影響。

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

3.19 日 Meetup 預告

北京時間本週六（3.19）上午 10:00，MegEngine Meetup 圍繞“Large Kernel Makes CNN Great Again”主題，將為大家帶來精彩線上分享。
活動資訊：週六直播預告 | 打破思維慣性，曠視MegEngine告訴你為什麼要思考大kernel size
直播間地址:https://live.bilibili.com/22436423

附：測試程式碼

MegEngine 測試程式碼

import time
import megengine.module as M
import megengine.autodiff as ad
import megengine
import numpy as np

megengine.functional.debug_param.set_execution_strategy("PROFILE")

def benchmark_lknet(ksize, batch=64, dim=384, res=32, depth=24):
    m = M.Sequential(
        *[M.Conv2d(dim, dim, ksize, padding=ksize//2, groups=dim, bias=False) for _ in range(depth)]
    )
    x = megengine.tensor(np.ones([batch, dim, res, res]))
    gm = ad.GradManager().attach(m.parameters())
    for i in range(20):
        t = time.perf_counter()
        with gm:
            y = m(x)
            gm.backward(y.mean())
        megengine._full_sync()
        t = time.perf_counter() - t
        if i > 9 and i % 10 == 0:
            print(t)
    return t

if __name__ == "__main__":
    args = dict()
    for k in (3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31):
        t = benchmark_lknet(k, **args)
        print("kernel size", k, "iter time", t * 1000, "ms")

PyTorch 測試程式碼

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import time

cudnn.benchmark = True

def benchmark_lknet(ksize, batch=64, dim=384, res=32, depth=24):
    m = nn.Sequential(
        *[nn.Conv2d(dim, dim, ksize, padding=ksize//2, groups=dim, bias=False) for _ in range(depth)]
    ).cuda()
    x = torch.rand(batch, dim, res, res).cuda()
    for i in range(20):
        t = time.perf_counter()
        y = m(x)
        y.mean().backward()
        torch.cuda.synchronize()
        t = time.perf_counter() - t
        if i > 9 and i % 10 == 0:
            print(t)
    return t

if __name__ == "__main__":
    args = dict()
    for k in (3, 5, 7, 9, 11, 13, 15, 17, 19, 21, 23, 25, 27, 29, 31):
        t = benchmark_lknet(k, **args)
        print("kernel size", k, "iter time", t * 1000, "ms")

GitHub：MegEngine 天元

官網：MegEngine-深度學習，簡單開發