作者：于雄雄 陳其友 | 曠視 MegEngine 架構(gòu)師

背景

在 CV 領(lǐng)域中，卷積計(jì)算是擴(kuò)充像素的感受野的有效方法，模型大多數(shù)的計(jì)算量都是卷積操作貢獻(xiàn)的。因此在 CV 模型的推理性能優(yōu)化中，最重要的一項(xiàng)工作是對(duì)卷積的優(yōu)化。MegEngine 在長期的工業(yè)界實(shí)踐和反饋的基礎(chǔ)上總結(jié)得出卷積優(yōu)化的基本方法有：

• 直接卷積計(jì)算優(yōu)化

該方法的計(jì)算過程為逐通道進(jìn)行卷積滑窗計(jì)算并累加，該優(yōu)化方法對(duì)卷積的參數(shù)敏感，為了達(dá)到最優(yōu)的性能，會(huì)根據(jù)各個(gè)卷積參數(shù)分別進(jìn)行 kernel 優(yōu)化，通用性弱，但是在 Depthwise 的卷積中卻是最高效的方法。

• FFT 卷積計(jì)算優(yōu)化

根據(jù)卷積的性質(zhì)，利用傅立葉變換可以將卷積轉(zhuǎn)換為頻域上的乘法，在頻域上的計(jì)算對(duì)應(yīng)乘法，再使用傅立葉變換逆變換，就可以得到卷積對(duì)應(yīng)的計(jì)算結(jié)果。該方法使用高性能的傅立葉變換算法，如 FFT，可以實(shí)現(xiàn)卷積計(jì)算的優(yōu)化，算法性能完全取決于傅立葉變換的性能以及相應(yīng)卷積參數(shù)。

• Im2col+matmul 卷積計(jì)算優(yōu)化

由于卷積計(jì)算中有大量的乘加運(yùn)算，和矩陣乘具有很多相似的特點(diǎn)，因此該方法使用 Im2col 的操作將卷積運(yùn)算轉(zhuǎn)化為矩陣運(yùn)算，最后調(diào)用高性能的 Matmul 進(jìn)行計(jì)算。該方法適應(yīng)性強(qiáng)，支持各種卷積參數(shù)的優(yōu)化，在通道數(shù)稍大的卷積中性能基本與 Matmul 持平，并且可以與其他優(yōu)化方法形成互補(bǔ)。

• Winograd 卷積計(jì)算優(yōu)化

Winograd 方法是按照 Winograd 算法的原理將卷積運(yùn)行進(jìn)行轉(zhuǎn)變，達(dá)到減少卷積運(yùn)算中乘法的計(jì)算總量。其主要是通過將卷積中的乘法使用加法來替換，并把一部分替換出來的加法放到 weight 的提前處理中，從而達(dá)到加速卷積計(jì)算的目的。Winograd 算法的優(yōu)化局限為在一些特定的常用卷積參數(shù)才支持。

由于 direct 卷積可以直接由公式得來，而 FFT 卷積對(duì)于當(dāng)前業(yè)界用到的各種參數(shù)的卷積，其性能優(yōu)勢遠(yuǎn)沒有其他優(yōu)化方法明顯，對(duì)于這兩者本文不做詳細(xì)展開。這里主要講述 Im2col 和 Winograd 算法的實(shí)現(xiàn)以及優(yōu)化方法。

Im2col+Matmul 優(yōu)化

Im2col 算法簡介

Im2col+Matmul 方法主要包括兩個(gè)步驟：

• 使用 Im2col 按照卷積核的需要將輸入矩陣展開一個(gè)大的矩陣，矩陣的每一列表示卷積核需要的一個(gè)輸入數(shù)據(jù)。

• 使用上面轉(zhuǎn)換的矩陣進(jìn)行 Matmul 運(yùn)算，得到的數(shù)據(jù)就是最終卷積計(jì)算的結(jié)果。

具體 Im2col 的步驟如上圖所示：

1. 將輸入數(shù)據(jù)按照卷積窗進(jìn)行展開并存儲(chǔ)在矩陣的列中，多個(gè)輸入通道的對(duì)應(yīng)的窗展開之后將拼接成最終輸出 Matrix 的一列。

2. 以卷積的 stride 為步長展開后續(xù)的卷積窗并存在 Matrix 的下一列中。

完成 im2col 的操作之后會(huì)得到一個(gè)輸入矩陣，卷積的 weights 也可以轉(zhuǎn)換為一個(gè)矩陣，此時(shí)對(duì) weights 的矩陣和 Im2col 的輸出矩陣進(jìn)行 Matmul 計(jì)算，就可以得到最終的卷積計(jì)算結(jié)果。

算法優(yōu)化

上面介紹的過程是原始 Im2col+Matmul 的過程，實(shí)際處理器在執(zhí)行上面的過程中性能達(dá)不到最優(yōu)，以輸入 Tensor 的 shape 為 (1, IC, IH, IW)，weights 的 shape 為 (OC，IC，F(xiàn)h，F(xiàn)w)，輸出 Tensor 的 shape 為 (1, OC, OH, OW) 為例，主要原因在于：

• Im2col 的輸入 Tensor 需要的 CPU 內(nèi)存大小為 IC*IH*IW，而按照上面 Im2col 之后所需要的內(nèi)存大小為 IC*Fh*Fw*OH*OW，當(dāng)卷積的 stride=1 的時(shí)候，Im2col 之后需要的內(nèi)存比之前大很多。

• 由于 Im2col 之后的數(shù)據(jù)量比較大，難以全部保存在 CPU 的 Cache 中，造成后續(xù) Matmul 計(jì)算時(shí)，讀取數(shù)據(jù)會(huì)存在 Cache Miss。

• Im2col 過程中會(huì)將輸入進(jìn)行 relayout 操作，而在后續(xù) Matmul 的計(jì)算中，需要對(duì)該數(shù)據(jù)進(jìn)行 Pack，Pack 操作會(huì)引入非必要的讀寫過程。影響算法實(shí)際性能。

優(yōu)化 1：對(duì) Im2col+Matmul 過程進(jìn)行分塊

上面提到在 Im2col 之后，消耗的內(nèi)存會(huì)超過 CPU 的 Cache 的容量，為了使這部分?jǐn)?shù)據(jù)能夠保存在 Cache 中，需要對(duì) Im2col+Matmul 的整個(gè)過程進(jìn)行分塊，每次 Im2col+Matmul 都只對(duì)一個(gè)分塊進(jìn)行操作，這樣就可以解決內(nèi)存占用過大，超過 CPU Cache 后造成 Cache Miss 的問題。

分塊優(yōu)化如上圖所示：Im2col 每次只對(duì) block_size 大小的數(shù)據(jù)進(jìn)行計(jì)算，得到的 Fh*Fw*IC*block_size 的數(shù)據(jù)可以保存在 Cache 中。Im2col 得到數(shù)據(jù)后，對(duì)其直接進(jìn)行 Matmul 計(jì)算，將計(jì)算得到的結(jié)果寫入到輸出 Tensor 對(duì)應(yīng)的 block_size 處就可以得到該分塊處卷積的計(jì)算結(jié)果。計(jì)算完該分塊之后，依次進(jìn)行下一個(gè) block_size 的計(jì)算，直到整個(gè)輸入計(jì)算完成。

結(jié)合?Matmul 的相關(guān)優(yōu)化知識(shí)，在進(jìn)行 Matmul A*B=C 計(jì)算時(shí)將分塊 Im2col 得到的數(shù)據(jù)視作 B 矩陣，A 矩陣為卷積的權(quán)重矩陣，根據(jù) sgemm 的分塊規(guī)則，以及 cache 的性質(zhì)，A 矩陣會(huì)被調(diào)度并保存在 L2 上，B 矩陣基于最內(nèi)層分塊的一列和 A 矩陣基于最內(nèi)存分塊的一行以及 C 矩陣基于最內(nèi)層的部分分塊會(huì)被調(diào)度保存在 L1 上，因此可以通過 L1，L2 的大小以及 A 矩陣的大小，計(jì)算出所有的分塊大小。下面是分塊優(yōu)化性能的試驗(yàn)結(jié)果，可以看出分塊優(yōu)化能有效的減少存儲(chǔ)使用，而且還可以提升算子的計(jì)算性能。

優(yōu)化 2：融合 Im2col 和 Matmul PACK 數(shù)據(jù)操作部分

Im2col 過程中將多個(gè)窗的展開同時(shí)進(jìn)行時(shí)，實(shí)際上是對(duì)內(nèi)存的 copy 以及數(shù)據(jù)的 relayout 的過程，后續(xù) Matmul 的 Pack 操作業(yè)是對(duì)數(shù)據(jù)的 copy 的 relayout，因此可以將上面兩次數(shù)據(jù)的 copy 和 relayout 進(jìn)行合并優(yōu)化，減少該過程中對(duì)內(nèi)存的讀寫次數(shù)。

如上圖所示 Im2col+Matmul 的 algo 中實(shí)現(xiàn)了將 Im2col 和 Matmul 的 Pack 融合的優(yōu)化，這樣能夠減少一次數(shù)據(jù)的讀寫操作。由于該 fuse 過程和卷積的參數(shù)直接相關(guān)，不同的卷積參數(shù)將對(duì)應(yīng)不同的融合 kernel，所以不具備通用性。通用情況下我們會(huì)使用之前的 Im2col+Matmul 的做法，另外針對(duì)一些通用的卷積如：kernel=3x3，stride=2 等，因?yàn)閰?shù)固定，因此可以直接進(jìn)行上述融合優(yōu)化，利用這樣的組合優(yōu)化，既可以保證 im2col 算法的通用性，也可以確保大部分常見的卷積的性能。

對(duì)融合之后的卷積進(jìn)行性能測試，如下所示為對(duì)應(yīng)的計(jì)算吞吐：

可以看出，大多數(shù)情況下，融合之后卷積會(huì)有明顯的性能提升。

Winograd 優(yōu)化

Winograd 算法簡介

Winograd 算法能夠優(yōu)化卷積計(jì)算的乘法計(jì)算量，乘法計(jì)算量的優(yōu)化原理可以參考相關(guān)論文。在此就不做過多介紹了。雖然 Winograd 可以優(yōu)化乘法的計(jì)算量，但是會(huì)增加加法的計(jì)算量，優(yōu)化這些加法的存在可以進(jìn)一步提高 Winograd 算法的性能。如可以把一部分加法計(jì)算提前到 weights 的預(yù)處理中，可以把部分加法隱藏在 Winograd 預(yù)處理中的 relayout 中。類似這樣的優(yōu)化可以達(dá)到減少卷積計(jì)算量的目的。

如下圖所示為 Winograd 卷積算法的基本步驟，主要包括：

• 把輸入的 feature map 和 weight 進(jìn)行 Winograd 轉(zhuǎn)換；

• 把轉(zhuǎn)換后 feature map 和 weight 做批量 Matmul；

• 把矩陣乘的結(jié)果進(jìn)行輸出轉(zhuǎn)換，得到最終結(jié)果。

在這些主要步驟中，要如何進(jìn)行 Winograd 轉(zhuǎn)換，如何 relayout，以及如何進(jìn)行輸出轉(zhuǎn)換呢？下面以 Winograd F(2x2, 3x3) 為例，詳細(xì)說明下這些過程。

如上圖所示，上半部分是 weights 的轉(zhuǎn)換，下半部分是輸入 FeatureMap 的轉(zhuǎn)換。其中包括了 Winograd 轉(zhuǎn)換以及 relayout 的過程。

對(duì)于 weights 的轉(zhuǎn)換，首先通過 Winograd 變換矩陣 G 和 GT 分別將 3x3 的 weight 轉(zhuǎn)換為 4x4 的矩陣，然后將該矩陣中相同位置的點(diǎn)（如圖中藍(lán)色為位置 1 的點(diǎn)）relayout 為一個(gè) IC*OC 的矩陣，最終形成 4x4=16 個(gè)轉(zhuǎn)換之后 weights 矩陣。

對(duì)于 FeatureMap 的轉(zhuǎn)換，首先將輸入 FeatureMap 按照 4x4 tile 進(jìn)行切分，然后將每個(gè) tile 通過 B 和 BT 轉(zhuǎn)換為 4x4 的矩陣，矩陣 B 和 BT 為 FeatureMap 對(duì)應(yīng)的 Winograd 變換矩陣，然后進(jìn)行與 weight 處理相似的 relayout，轉(zhuǎn)換為 16 個(gè) nr_tiles*IC 的 FeatureMap 矩陣。

如上圖所示，將上述轉(zhuǎn)換后兩批矩陣做矩陣乘，得到 16 個(gè) nr_tiles*OC 的矩陣，然后將相同位置的 16 個(gè)點(diǎn)轉(zhuǎn)換為 nr_tiles*OC 個(gè) 4x4 矩陣，再使用輸出的 Winograd 變換矩陣 A 和 AT 將這些 4x4 的矩陣轉(zhuǎn)換為 2x2 的輸出矩陣，最后將這些矩陣寫回輸出矩陣中就可以得到 Winograd 卷積的最終結(jié)果。

算法優(yōu)化

優(yōu)化 1：weight 提前處理

在上述 Winograd 算法的基礎(chǔ)上，鑒于模型中的權(quán)重?cái)?shù)據(jù)在整個(gè) Inference 的時(shí)候已經(jīng)是常量不會(huì)再改變，因此可以在真正 Inference 之前就可以對(duì)模型進(jìn)行了 weights 的轉(zhuǎn)換，這樣可以優(yōu)化在 Inference 的時(shí)候 weights 轉(zhuǎn)換的開銷，特別是在 IC 和 OC 較大時(shí)，weight 轉(zhuǎn)換的開銷非常大，所以 weights 提前轉(zhuǎn)換，尤其對(duì) Winograd 優(yōu)化特別重要，下圖是 Winograd 中進(jìn)行 weight 提前轉(zhuǎn)換和不進(jìn)行 weight 提前轉(zhuǎn)換時(shí)各自的性能：

從上圖可以看出 weight 轉(zhuǎn)換在 Winograd 中耗時(shí)占比很大，進(jìn)行 weight 提前轉(zhuǎn)換可以帶來很大的性能收益。

優(yōu)化 2：Winograd 分塊優(yōu)化

上述的 Winograd 算法，還會(huì)有以下缺點(diǎn)：

1. 輸入轉(zhuǎn)換需要跨 channel 讀寫整個(gè) feature map，數(shù)據(jù)讀寫對(duì) Cache 不友好。

2. feature map 轉(zhuǎn)換之后，矩陣乘時(shí)需要再 PACK，數(shù)據(jù)訪存增加。

針對(duì)這些問題，可以對(duì) Winograd 算法的整個(gè)計(jì)算流程做進(jìn)一步的優(yōu)化，這些優(yōu)化主要包括：

• 輸入轉(zhuǎn)換時(shí)，分塊 feature map 的 tiles 進(jìn)行分塊，每次只進(jìn)行一定數(shù)量的 tiles 計(jì)算；

• 調(diào)整分塊大小適配 CPU L1 Cache，使得矩陣乘不需要 PACK；

對(duì)整個(gè)輸入 feature map 進(jìn)行分塊后，每次只計(jì)算一個(gè)分塊的 nr 個(gè) tiles，這樣就可以保證每個(gè)批量矩陣的輸入數(shù)據(jù)（除了轉(zhuǎn)換之后的 weight 數(shù)據(jù)）保存于 L1 Cache，不會(huì)出現(xiàn) Cache miss， 而且矩陣乘時(shí)不需要 PACK。

下面是分塊優(yōu)化前后的速度對(duì)比，可以看出分塊優(yōu)化對(duì)性能有顯著的提升。

總結(jié)

CPU 上 Inference 中有關(guān)卷積的優(yōu)化有很多的途徑，這里我們主要介紹了 Im2col+matmul 卷積以及 Winograd 卷積中的一些進(jìn)一步優(yōu)化的技術(shù)手段，通過這些方法可以進(jìn)一步加速卷積計(jì)算的性能，從而加速整個(gè)模型的 Inference 性能。如下圖所示是 float32 的經(jīng)典網(wǎng)絡(luò)開啟相關(guān)優(yōu)化后，在驍龍 855 上的測試速度：

對(duì)于具體的優(yōu)化細(xì)節(jié)，大家可以結(jié)合 MegEngine 的代碼實(shí)現(xiàn)進(jìn)行研究，歡迎大家提出寶貴意見。

更多 MegEngine 信息獲取，您可以查看：

文檔：https://www.megengine.org.cn/doc/stable/zh/?

深度學(xué)習(xí)框架 MegEngine 官網(wǎng)：https://www.megengine.org.cn/

GitHub 項(xiàng)目：https://github.com/MegEngine，或加入 MegEngine 用戶交流 QQ 群：1029741705