在 CFU Playground 上加速 MLPerf™ Tiny 影像分類模型 #2:直接卷積與 im2col

Yeecy
17 min readJan 8, 2024

--

2.1 直接卷積

卷積層有許多參數,有卷積核(kernel)的尺寸和數量、有卷積時的步長(stride)、填充值(padding)、有卷積的擴張值(dilation)等等,為了文章的簡潔,在此先約定一些符號來代稱,未列出的參數就表示本文不會討論(例如 group)。

  • 輸入高度 Hᵢₙ
  • 輸入寬度 Wᵢₙ
  • 輸入通道數 Cᵢₙ
  • 卷積核高度 H
  • 卷積核寬度 W
  • 輸出高度 Hₒᵤₜ
  • 輸出寬度 Wₒᵤₜ
  • 輸出通道數 Cₒᵤₜ
  • 高度步長值 H
  • 寬度步長值 W
  • 填充高度 H
  • 填充寬度 W
  • 高度擴張值 Hₑ(因為 Unicode 沒有下標 d,所以用 e 代替)
  • 寬度擴張值 W

相信會點進這篇文章的讀者對卷積層的運作原理有一定的了解,因此接下來只會快速複習一下重點部分,不會太深入講解,我們的重點會放在下一小節 im2col 上。

在一張彩色圖片中會有三個通道分別表示紅綠藍三色的值,若我們不做其他處理直接將其餵給卷積層,假設該圖的大小為 128×64,則對於該卷積層而言 Hᵢₙ 為 128,Wᵢₙ 為 64,Cᵢₙ 為 3,也就是說這三個參數對於卷積層而言並不固定,至於其他除了 Hₒᵤₜ 和 Wₒᵤₜ 以外的參數則可以被人們自由調整。

對於 Hₒᵤₜ 而言,其受到 Hᵢₙ、Hₚ、Hₛ 和 Hₑ 影響,關係式為 Hₒᵤₜ = floor((Hᵢₙ + 2Hₚ - HHₑ + Hₑ - 1) ÷ Hₛ) + 1,Wₒᵤₜ 依此類推。在不同的需求下,我們可以透過上面的關係式推得不同的參數值,舉例來說,如果希望 Hₒᵤₜ 與 Hᵢₙ 相等,我們可以透過固定 Hₖ、Hₛ 和 Hₑ 後透過關係式求 Hₚ。

這裡用一些小例子觀察不同參數的影響,先設 Hᵢₙ = Wᵢₙ = 3、Cᵢₙ = Cₒᵤₜ = 1、Hₖ = Wₖ = 2 做為這些例子的共同參數,輸入和卷積核約定如下,其中 [] 表示一個數值,另外令超出邊界的值為 0,且不考慮偏差值。

input
-----------
[A] [B] [C]
[D] [E] [F]
[G] [H] [I]

kernel
-------
[w] [x]
[y] [z]

情況一:Hₚ = Wₚ = 0、Hₛ **= Wₛ = 1、H= Wₑ = 1

因為 Hₒᵤₜ = Wₒᵤₜ = floor((3 + 2×0 - 2×1 + 1 - 1) ÷ 1) + 1 = 2,所以輸出共有四個元素。

output
---------------------------
[Aw+Bx+Dy+Ez] [Bw+Cx+Ey+Fz]
[Dw+Ex+Gy+Hz] [Ew+Fx+Hy+Iz]

情況二:Hₚ = Wₚ = 1、Hₛ **= Wₛ = 1、H= Wₑ = 1

因為 Hₒᵤₜ = Wₒᵤₜ = floor((3 + 2×1 - 2×1 + 1 - 1) ÷ 1) + 1 = 4,所以輸出共有 16 個元素。

output
-------------------------------------------
[Az] [Ay+Bz] [By+Cz] [Cy]
[Ax+Dz] [Aw+Bx+Dy+Ez] [Bw+Cx+Ey+Fz] [Cw+Fy]
[Dx+Gy] [Dw+Ex+Gy+Hz] [Ew+Fx+Hy+Hz] [Fw+Iy]
[Gx] [Gw+Hx] [Hw+Ix] [Iw]

情況三:Hₚ = Wₚ = 0、Hₛ = Wₛ = 2、H= Wₑ = 1

因為 Hₒᵤₜ = Wₒᵤₜ = floor((3 + 2×0 - 2×1 + 1 - 1) ÷ 2) + 1 = 1,所以輸出只有一個元素。

output
-------------
[Aw+Bx+Dy+Ez]

情況四:Hₚ = Wₚ = 0、Hₛ = Wₛ = 1、H= Wₑ = 2

因為 Hₒᵤₜ = Wₒᵤₜ = floor((3 + 2×0 - 2×2 + 2 - 1) ÷ 1) + 1 = 1,所以輸出也只有一個元素。

output
-------------
[Aw+Cx+Gy+Iz]

從上面的例子可以看到當 Hₚ 大於 0 時,會導致計算卷積有機會存取到超出邊界的值,需要特別處理,這一點對於在軟體層面進行最佳化相當重要。

最後我們用個大一點的例子以便銜接下一小節,此處主要演示 Cᵢₙ 和 Cₒᵤₜ 大於 1 的情況。

Hᵢₙ = Wᵢₙ = 2、Cᵢₙ = Cₒᵤₜ = 2、Hₖ = Wₖ = 2、Hₚ = Wₚ = 1、Hₛ = Wₛ = 1、Hₑ = Wₑ = 1 且輸入和卷積核如下。

input
------------------
ch1 ch2
-------- --------
[A] [B] [E] [F]
[C] [D] [G] [H]

kernel
-------------------------------------
ch11 ch12 ch21 ch22
-------- ------- -------- -------
[k] [l] [o] [p] [s] [t] [w] [x]
[m] [n] [q] [r] [u] [v] [y] [z]

因為 Hₒᵤₜ = Wₒᵤₜ = floor((2 + 2×1 - 2×1 + 1 - 1) ÷ 1) + 1 = 3 且 Cₒᵤₜ 為 2,共會產生 18 個元素。

output
-----------------------------------------------------
ch1
-----------------------------------------------------
[An+Er] [Am+Bn+Eq+Fr] [Bm+Fq]
[Al+Cn+Ep+Gr] [Ak+Bl+Cm+Dn+Eo+Fp+Gq+Hr] [Bk+Dm+Fo+Hq]
[Cl+Gp] [Ck+Dl+Go+Hp] [Dk+Ho]

ch2
-----------------------------------------------------
[Av+Ez] [Au+Bv+Ey+Fz] [Bu+Fy]
[At+Cv+Ex+Gz] [As+Bt+Cu+Dv+Ew+Fx+Gy+Hz] [Bs+Du+Fw+Hy]
[Ct+Gx] [Cs+Dt+Gw+Hx] [Ds+Hw]

不知道讀者看到這裡,有沒有發現卷積層也是滿滿的乘積累加呢?

最後我們用卷積層的程式碼(這是 Python 喔)作為本小節的結尾,以下的程式碼還可以進行循環不變量外提(loop-invariant code motion)和迴圈順序交換(loop interchange),兩種方法都有可能帶來性能獲益,不過通常編譯器能很容易地幫我們做到前者,而後者需要人們自己手動調整。

# Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
# Direct Convolution

for ch_out in range(C_out):
for y_out in range(H_out):
for x_out in range(W_out):
acc = 0

for ch_in in range(C_in):
for y_k in range(H_k):
for x_k in range(W_k):
y_in = y_out*H_s + y_k*H_e - H_p
x_in = x_out*W_s + x_k*W_e - W_p
if 0 <= y_in < H_in and 0 <= x_in < W_in:
acc += data[y_in][x_in][ch_in] \\
* kernel[ch_out][y_k][x_k][ch_in]

output[y_out][x_out][ch_out] = acc

這段程式碼片段的核心為 y_inx_in 的計算,這邊以 y_in 舉例, x_in 可以類推適用。在表達式 y_in = y_out*H_s + y_k*H_e - H_p 中最重要的是 y_iny_out 的關係,首先 Hₛ 表示的是卷積核在垂直方向的移動步長,因此當 y_out 增加 1,y_in 需要增加 Hₛ,因此 y_out 自然需要乘上 H_s

接著考慮 Hₑ 和 Hₖ,當 Hₑ 為 1 時,表示卷積核實際對應到的輸入位置是連續的,也就是說當 y_k 增加 1,y_in 也增加 1,那如果 Hₑ 為 2 呢?Hₑ 為 2 表示卷積核上相鄰的兩個位置間有個空洞(等價於相鄰兩個值之間都有個不可學習的常數 0),因此當 y_k 增加 1 時,對應到的 y_in 需要增加 2,因為 y_in + 1 對應到了空洞(也可以說是 y_in + 1 乘上 0 後為 0,所以可以忽略不算),所以當 y_k 增加 Hₑ, y_in 也該增加 Hₑ,即 y_k 需要乘上 H_e

最後考慮 Hₚ,當 Hₚ 為 0 時,顯然 yᵢₙ 的取值範圍必然在 0 到 Hᵢₙ - 1 之間,不會產生任何問題,然而若 Hₚ 為 1,等價於我們在原先的輸入資料的每個通道的最上方和最下方個別多填充一排 0,如果輸入的資料在送入卷積層前就經過填充,則我們不需要在表達式中考慮 H_p,直接運算即可,但因為 Hᵢₙ 和 Hₒᵤₜ 通常很大,先進行填充會耗費大量時間與空間,所以在邏輯上使得 yᵢₙ = -1 和 yᵢₙ = Hᵢₙ 對應到 0 較為划算,那要怎麼做到呢?其實很簡單,把 y_out*H_s + y_k*H_e 減去 Hₚ 即可,例如當 y_outy_k 為 0 時(這個情況必然會發生),yᵢₙ 應該對應到最上面一排的某個值,也就是說 yᵢₙ 應該從原本的 0 對應到 -1,也就是將 0 減去 Hₚ = 1 得到 -1。那麼 yᵢₙ = Hᵢₙ 的情況會發生嗎?答案是會,因為我們在計算 Hₒᵤₜ 時已經加上了 2Hₚ,所以必然會遇到 yᵢₙ = Hᵢₙ 的情況。

2.2 Im2col

相信讀者聽聞過近年來 NVIDIA 在 GPU 裡新增了 Tensor 核心,雖然這個名字聽起來讓人覺得它能做張量運算,但實際上(至少從 Volta 到 Hopper 架構都如此)Tensor 核心本身一次計算只能完成一次固定大小的矩陣乘法,例如在 Volta 架構上一個週期可以完成一次兩個 4x4 矩陣的乘法,也就是說應該叫 Matrix 核心比較貼切。

總之,如果 Tensor 核心只能做矩陣乘法,NVIDIA 的計算卡到底要怎麼加速卷積神經網路呢?這個問題的思路簡單暴力,那就是直接把卷積層運算變換成等價的矩陣乘法!這個方法的名稱即為這個小節的名稱 im2col。

Im2col 的全稱為 image to column,也就是把輸入的圖片轉變為矩陣的行,但為了配合 im2col,我們實際上還需要對卷積核做相應的調整。

對於卷積層,若 batch size 為 1,則輸入資料的大小為 Hᵢₙ×Wᵢₙ×Cᵢₙ,卷積核的大小為 Cₒᵤₜ×Hₖ×Wₖ×Cᵢₙ,輸出結果的大小為 Hₒᵤₜ×Wₒᵤₜ×Cₒᵤₜ,如果要用矩陣乘法來實現卷積運算,那麼我們需要將輸入資料和卷積核輾平成二維的矩陣,並將計算完所得的二維矩陣還原回原本的三維輸出。

因為輸出的大小為 Hₒᵤₜ×Wₒᵤₜ×Cₒᵤₜ,不妨讓矩陣乘法的結果大小為 HₒᵤₜWₒᵤₜ×Cₒᵤₜ,那麼 im2col 矩陣的第一個維度需為 HₒᵤₜWₒᵤₜ,卷積核矩陣的第二個維度需為 Cₒᵤₜ,這樣兩者相乘的結果才會是 HₒᵤₜWₒᵤₜ×Cₒᵤₜ。接下來的問題是剩下的維度應該要是多少,在卷積層中,卷積核上的每個值都會被用到,所以自然而然卷積核矩陣的維度為 CᵢₙHWₖ×Cₒᵤₜ,也就是說 im2col 矩陣的維度應為 HₒᵤₜWₒᵤₜ×CᵢₙHWₖ,如此一來一個 HₒᵤₜWₒᵤₜ×CᵢₙHWₖ 的矩陣乘上一個 CᵢₙHWₖ×Cₒᵤₜ 的矩陣就能得到一個 HₒᵤₜWₒᵤₜ×Cₒᵤₜ 的矩陣了!

我們將上一小節的最後一個範例以矩陣乘法的形式再做一遍,方便讀者比較兩者差異。

input
------------------
ch1 ch2
-------- --------
[A] [B] [E] [F]
[C] [D] [G] [H]

kernel
-------------------------------------
ch11 ch12 ch21 ch22
-------- ------- -------- -------
[k] [l] [o] [p] [s] [t] [w] [x]
[m] [n] [q] [r] [u] [v] [y] [z]

首先從上一小節可以知道 im2col 矩陣的大小為 9×8 而卷積核矩陣的大小為 8×2,兩個矩陣的構造方法會在稍後給出,讀者可以先根據下面的結果找找規律。

im2col
-------------------------------
[0] [0] [0] [A] [0] [0] [0] [E]
[0] [0] [A] [B] [0] [0] [E] [F]
[0] [0] [B] [0] [0] [0] [F] [0]
[0] [A] [0] [C] [0] [E] [0] [G]
[A] [B] [C] [D] [E] [F] [G] [H]
[B] [0] [D] [0] [F] [0] [H] [0]
[0] [C] [0] [0] [0] [G] [0] [0]
[C] [D] [0] [0] [G] [H] [0] [0]
[D] [0] [0] [0] [H] [0] [0] [0]

kernel
-------
[k] [s]
[l] [t]
[m] [u]
[n] [v]
[o] [w]
[p] [x]
[q] [y]
[r] [z]

im2col x kernel
---------------------------------------------------
[An+Er] [Av+Ez]
[Am+Bn+Eq+Fr] [Au+Bv+Ey+Fz]
[Bm+Fq] [Bu+Fy]
[Al+Cn+Ep+Gr] [At+Cv+Ex+Gz]
[Ak+Bl+Cm+Dn+Eo+Fp+Gq+Hr] [As+Bt+Cu+Dv+Ew+Fx+Gy+Hz]
[Bk+Dm+Fo+Hq] [Bs+Du+Fw+Hy]
[Cl+Gp] [Ct+Gx]
[Ck+Dl+Go+Hp] [Cs+Dt+Gw+Hx]
[Dk+Ho] [Ds+Hw]

有些讀者可能會感到疑惑,上面的方法看起來明明是 im2row 啊,為什麼要叫 im2col 呢?呵呵,筆者也不知道,問就是約定俗成,反正 im2row 和 im2col 本質上沒什麼區別,掌握了這裡介紹的方法後,想怎麼改都可以。總之將矩陣乘法的結果重排後,可以得到最後的輸出,讀者可以比較一下兩種方法得到的輸出結果是否相同。

output
-----------------------------------------------------
ch1
-----------------------------------------------------
[An+Er] [Am+Bn+Eq+Fr] [Bm+Fq]
[Al+Cn+Ep+Gr] [Ak+Bl+Cm+Dn+Eo+Fp+Gq+Hr] [Bk+Dm+Fo+Hq]
[Cl+Gp] [Ck+Dl+Go+Hp] [Dk+Ho]

ch2
-----------------------------------------------------
[Av+Ez] [Au+Bv+Ey+Fz] [Bu+Fy]
[At+Cv+Ex+Gz] [As+Bt+Cu+Dv+Ew+Fx+Gy+Hz] [Bs+Du+Fw+Hy]
[Ct+Gx] [Cs+Dt+Gw+Hx] [Ds+Hw]

2.2.1 構造卷積核矩陣

構造卷積核矩陣非常簡單,基本上就只是按照通道把卷積核輾平而已,對應到同一輸出通道的卷積核按輸入通道順序排列成一行,讀者應該能很容易地從上面的範例中看出這點。

# Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
# Kernel Matrix Construction

for ch_in in range(C_in):
for y_k in range(H_k):
for x_k in range(W_k):
for ch_out in range(C_out):
i = ch_in*H_k*W_k + y_k*W_k + x_k
j = ch_out
matrix_kernel[i][j] = kernel[ch_out][y_k][x_k][ch_in]

程式碼中的 i = ch_in*H_k*W_k + y_k*W_k + x_k 只是對於 CᵢₙHWₖ 的陣列索引值計算,對於修過 C/C++ 課程算過多維陣列和一維陣列間轉換的讀者應該很簡單,筆者就不多解釋了。

2.2.2 構造 im2col 矩陣

構造 im2col 矩陣相對麻煩一些,因為我們在算 Hₒᵤₜ 和 Wₒᵤₜ 時用到的參數通通都會出現在我們構造矩陣的程式碼裡,如果讀者不熟卷積運算的話,可以翻回前面慢慢思索,此處的重點依舊為前面卷積程式碼出現的 y_in = y_out*H_s + y_k*H_e - H_p,理解該表達式且能看懂卷積核矩陣的構造,那麼構造 im2col 矩陣自然難不倒讀者。

在前面的卷積核矩陣構造中,我們將同一輸出通道的卷積核按行從上至下展開排列,如果要讓 im2col 矩陣右乘卷積核矩陣得到卷積結果,我們需要想辦法將卷積核對應到的輸入值從左至右展開填充進 im2col 矩陣中,所以在下面的程式碼中需要分別計算跟卷積核和輸入資料相關的索引值。

# Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
# Im2col Matrix Construction in Im2row Manner

for y_out in range(H_out):
for x_out in range(W_out):
for ch_in in range(C_in):
for y_k in range(H_k):
for x_k in range(W_k):
i = y_out*W_out + x_out
j = ch_in*H_k*W_k + y_k*W_k + x_k
y_in = y_out*H_s + y_k*H_e - H_p
x_in = x_out*W_s + x_k*W_e - W_p

if 0 <= y_in < H_in and 0 <= x_in < W_in:
matrix_im2col[i][j] = data[y_in][x_in][ch_in]
else:
matrix_im2col[i][j] = 0

注意到當 Hₚ 大於 0 時有機會讀到圖片外面的值,因此需要額外邏輯判斷是否超出範圍,此處假設填充值為 0。

2.2.3 還原輸出結果

矩陣乘法我們會在下一節討論,這裡只討論如何把 HₒᵤₜWₒᵤₜ×Cₒᵤₜ 的乘法結果還原回原本輸出的形狀 Hₒᵤₜ×Wₒᵤₜ×Cₒᵤₜ。

# Copyright (c) 2023-2024 Chung-Yi Chen (Yeecy)
# Output Reconstruction

for y_out in range(H_out):
for x_out in range(W_out):
for ch_out in range(C_out):
i = y_out*W_out + x_out
j = ch_out
output[y_out][x_out][ch_out] = matrix_result[i][j]

觀察前面的範例,讀者應該也能很容易理解上面的程式碼。

最後在結束 im2col 的介紹之前,讓我們思考一下用矩陣乘法替代直接卷積運算的影響。

  • 計算量會降低嗎?
    — 不會。若兩種方法對任意輸入資料和卷積核所得的的運算結果相同,表示兩者需要的乘積累加操作數和記憶體存取次數相同,因此計算量不會減少,如果加上構造矩陣的開銷,那麼計算量會比直接卷積還高。
  • 記憶體使用量會增加嗎?
    如果要構造出 im2col 矩陣,我們知道該矩陣共有 HₒᵤₜWₒᵤₜCᵢₙHWₖ 個元素,而原本的輸入資料僅有 CᵢₙHᵢₙWᵢₙ 個元素,一般而言可以預期 HₒᵤₜWₒᵤₜHWₖ 大於 HᵢₙWᵢₙ,所以不僅需要額外記憶體空間,而且還是比原本輸入資料還大的記憶體空間,因此記憶體使用量會增加。
  • 用矩陣乘法的優點是什麼?
    — 循環層(recurrent layer)和注意力機制(attention)本質上都是矩陣乘法,如果把卷積層也變成矩陣乘法,那麼我們只需針對矩陣乘法進行最佳化,便能加速絕大多數的模型。
    — 硬體供應商通常會提供精調過的線性代數函式庫(Basic Linear Algebra Subprograms,BLAS)以便進行矩陣乘法,例如 Intel 的 MKL 和 NVIDIA 的 cuBLAS。
    — 可以應用軟體層面的矩陣乘法最佳化技巧或甚至是更改演算法(如 Strassen 演算法)
    — 可以在硬體上實現專門用來計算矩陣乘法的電路

總結來說,我們可以期待矩陣乘法的加速收益會比 im2col 的開銷來得高,因此總體而言矩陣乘法會比直接卷積來得快。

--

--

Yeecy

A Ph.D. student at NYCU CS and a compiler engineer at ICEshell Co., Ltd. You can find more information about me on my GitHub page github.com/ADNRs.