基于OneFlow實現(xiàn)Unfold Fold算子

從卷積層說起

熟悉CNN的小伙伴應(yīng)該知道卷積是一個很常用也很重要的操作，CNN里的卷積和信號處理的卷積并不是一回事，CNN的卷積是做一種二維的互相關(guān)運算，以《動手學(xué)深度學(xué)習(xí)》5.1章為示例：

窗口內(nèi)的元素和卷積核相乘，求和得到輸出元素，一份naive的代碼如下（同來自《動手學(xué)深度學(xué)習(xí)》

from?mxnet?import?autograd,?nd
from?mxnet.gluon?import?nn

def?corr2d(X,?K):??#?本函數(shù)已保存在d2lzh包中方便以后使用
????h,?w?=?K.shape
????Y?=?nd.zeros((X.shape[0]?-?h?+?1,?X.shape[1]?-?w?+?1))
????for?i?in?range(Y.shape[0]):
????????for?j?in?range(Y.shape[1]):
????????????Y[i,?j]?=?(X[i:?i?+?h,?j:?j?+?w]?*?K).sum()
????return?Y

這里是借助了numpy array的索引特性來寫的，如果在c++里寫，需要的循環(huán)層數(shù)更多（會在后面進行展示）。而這種循環(huán)計算的寫法效率并不高，會極大拖慢卷積運算的速度。

初見img2col

為了提高卷積運算的速度，img2col算法被發(fā)明了出來的，它的本質(zhì)是用矩陣乘法來等價卷積運算，示例圖如下：

https://github.com/microsoft/AI-System/blob/main/docs/SystemforAI-4-Computer%20architecture%20for%20Matrix%20computation.pdf 這是微軟的AISystem倉庫對應(yīng)的章節(jié)，強烈推薦大家去學(xué)習(xí)（本人鴿了好久沒看完）

可以看到img2col把輸入特征圖進一步展開，然后把filter展平，兩者做矩陣乘，得到了相同的結(jié)果。

理解img2col

看了上面的圖可能還是不能理解這是怎么展開的，這里我們會介紹下：

• 假設(shè)輸入特征圖為(N, Cin, H, W)，卷積核參數(shù)為(K, K, Cin, Cout), 輸出特征圖的長寬為(Oh, Ow)

經(jīng)過img2col后，輸入特征圖會變換為 (N, Cin*K*K, Oh*Ow) 這么一個三維向量。

此外卷積核我們會reshape成一個二維向量：(Cout, K*K*Cin)。

而這兩個向量可以做一個矩陣乘，輸出向量為(N, Cout, Oh*Ow)，這也是我們預(yù)期的卷積結(jié)果。

img2col算法是一種用空間換時間的做法，將輸入進行變換后，顯然它所占用的顯存空間變大了，好處則是可以借助矩陣乘，快速地完成卷積運算。

下面我會結(jié)合darknet的原生img2col和一篇博客來進行相關(guān)講解

代碼：https://github.com/pjreddie/darknet/blob/master/src/im2col.c

博客：https://blog.csdn.net/caicaiatnbu/article/details/100515321

img2col源碼

darknet的img2col其實是搬運自caffe的，我們這里為了方便理解，以簡單的CPU版本為例子，且不考慮batch維度。

為了讓讀者能夠快速運行上代碼，這里講解我以自己寫的一版darknet img2col來作為示例。

def?darknet_img2col(data,?channels,?height,?width,?ksize,?stride,?pad):
????out_h?=?int((height?+?2*pad?-?ksize)?/?stride)?+?1
????out_w?=?int((width?+?2*pad?-?ksize)?/?stride)?+?1

????channels_cols?=?channels*ksize*ksize
????out_shape?=?(channels_cols,?out_h*out_w)
????elem_cnt?=?out_shape[0]?*?out_shape[1]

????out_array?=?np.zeros(shape=elem_cnt,?dtype=np.float32)

首先我們可以確定輸出tensor的各個維度，其中out_h和out_w就是輸出的高，寬，采用的是卷積輸出的公式：

而channel_cols則是之前我們提到的，img2col會把第二個維度變換為C_in*K*K。

然后進入到次數(shù)為channel_cols的for循環(huán)

????for?c?in?range(channels_cols):
????????#?分別計算一個k*k的窗口內(nèi)，h，w的偏移量
????????kh_offset?=?(c?//?ksize)?%?ksize
????????kw_offset?=?c?%?ksize
????????#?計算當(dāng)前處理的通道index
????????c_im?=?c?//?ksize?//?ksize

然后我們需要根據(jù)當(dāng)前處理的輸出元素index，來獲取對應(yīng)輸入的元素

????????for?h?in?range(out_h):
????????????for?w?in?range(out_w):
????????????????im_row?=?kh_offset?+?h?*?stride
????????????????im_col?=?kw_offset?+?w?*?stride
????????????????index?=?(c?*?out_h?+?h)?*?out_w?+?w

im_row的計算方式邏輯是：當(dāng)前處理的輸入元素窗口起始點：即h*stride，然后加上窗口內(nèi)的kh_offset偏移量。

而index的計算比較容易，因為輸出是(C, H, W)，對應(yīng)的一維index那就是

當(dāng)前channel*Oh*Ow + 當(dāng)前h*Ow + 當(dāng)前w

最后我們再將元素取出來，賦給out數(shù)組。然后再將一維的out數(shù)組reshape成我們前面推導(dǎo)得到的out_shape

????????????????out_array[index]?=?im2col_get_pixel(data,?height,?width,?c_im,?im_row,?im_col,??pad)

????out_array?=?np.reshape(out_array,?out_shape)

????return?out_array

img2col_get_pixel是一個合法取元素的函數(shù)，如果出現(xiàn)越界范圍（比如小于0，或者大于Oh），那么就是取到padding的部分，此時我們返回0。

def?im2col_get_pixel(im,?height,?width,?channel,?row,?col,?pad):
????row?-=?pad
????col?-=?pad
????if?row?0?or?col?0?or?row?>=?height?or?col?>=?width:
????????return?0
????return?im[int(col?+?width?*?(row?+?height?*?channel))]?#?batch*w*h*c?+?width*height*channel?+?width*row?+?col

我們可以簡單構(gòu)造一個數(shù)組來驗證結(jié)果（以微軟AI-System課程的示例作為輸入）

x?=?np.arange(1,?10).astype(np.float32)
out?=?darknet_img2col(x,?channels=1,?height=3,?width=3,?ksize=2,?stride=1,?pad=0)

輸出結(jié)果符合預(yù)期：

[[1.?2.?4.?5.]
?[2.?3.?5.?6.]
?[4.?5.?7.?8.]
?[5.?6.?8.?9.]]

col2img

col2img則是img2col的逆過程，有興趣的讀者可以參考下這個博客：

https://blog.csdn.net/caicaiatnbu/article/details/102626135

在后面的oneflow實現(xiàn)里會有更完整的圖例用于理解

OneFlow對應(yīng)的實現(xiàn)

PR地址：https://github.com/Oneflow-Inc/oneflow/pull/5675

OneFlow版本的Unfold

在深度學(xué)習(xí)框架里，img2col和col2img在Pytorch里還有另外的名字，也就是Unfold和Fold。通常想自己自定義一些跟卷積相關(guān)的騷操作時候，就經(jīng)常會用到這兩個Op。

我們假設(shè)輸入是一個(1, 2, 4, 4)的張量，但在框架內(nèi)部，我們通常都是以一個一維數(shù)組來存儲的，如下圖所示：

然而我們需要對應(yīng)的高維數(shù)組索引，OneFlow內(nèi)部有一個NdIndexHelper類，構(gòu)造的時候我們可以傳入高維shape，然后調(diào)用OffsetToNdIndex來進行一維offset到高維索引的轉(zhuǎn)換。

這里我們分別給輸入，輸出Tensor構(gòu)造一個NdIndexHelper

in_index_helper?=?NdIndexOffsetHelper(input_dims);?//?格式為(N,?C,?H,?W)
out_index_helper?=?NdIndexOffsetHelper(output_dims);?//?格式為(N,?C,?Kh,?Kw,?Oh,?Ow)

比較特別的是，我們把輸出構(gòu)造成一個6維的形式。

接著就是從輸出的視角來推導(dǎo)應(yīng)該取哪一個輸入元素

//?index_a?format:?(N,?C,?di,?hi,?wi,?db,?hb,?wb)?or?(N,?di,?hi,?wi,?db,?hb,?wb,?C)
//?index_b?format:?(N,?C,?D,?H,?W)?or?(N,?D,?H,?W,?C)
//?return:?true?indicates?out-of-bound,?otherwise?in-bound
template<typename?INDEX_T,?int?NDIM,?int?SDIM>
OF_DEVICE_FUNC?bool?UnfoldIndexTransform(const?UnfoldParams&?params,
?????????????????????????????????????????const?INDEX_T*?index_a,?INDEX_T*?index_b)?{
??//?batch?dim?index?transform
??index_b[0]?=?index_a[0];
??//?channel?dim?index?transform
??using?ParamType?=?UnfoldParams;
??index_b[ParamType::kInputChannelDim]?=?index_a[ParamType::kOutputChannelDim];
//?spatial?dim?index?transform

??//?D,H,W?spatial?dim?index?transform
??for?(int64_t?d?=?0;?d?????INDEX_T?idx?=?index_a[SDIM?+?NDIM?+?d]?*?params.stride[d]
??????????????????+?index_a[SDIM?+?d]?*?params.dilation[d]?-?params.padding[d];
????if?(idx?0?||?idx?>=?params.dims[d])?return?true;
????index_b[SDIM?+?d]?=?idx;
??}
??return?false;
}

• 模板參數(shù) INDEX_T表示Index的數(shù)據(jù)類型（可以有int32_t, int64_t），NDIM表示處理幾維數(shù)據(jù)（這里我們是2維），SDIM則是決定通道維所在位置，SDIM=1是NHWC格式，SDIM=2則是NCHW格式（這里我們?nèi)?）
• 輸入?yún)?shù) index_a表示輸出的NdIndexHelper，index_b則表示的是輸入的NdIndexHelper

從前面我們可以看到N，C這兩個維度的index是不變的，所以我們直接給過去

然后進入一個次數(shù)為NDIM==2的循環(huán)

這里index的計算是從輸出往輸入推，公式是（以H為例）：

Oh*stride_h + kh*dilation_h - padding_h

計算得到輸入的index，如果小于0或者大于輸入的寬高，則說明是到了padding的地方，我們直接return true，以表示越界。如果能取到元素，我們則將這個index賦給index_b即輸入的NdIndexHelper，且return false。

這部分的分解操作可參考下圖：

從輸出推導(dǎo)的好處就是整個運算是一個elementwise的操作，我們可以用輸出tensor的元素個數(shù)做一個循環(huán)完成整個unfold操作。

template<typename?T,?typename?INDEX_T,?int?NDIM,?int?SDIM>
__global__?void?CudaUnfoldForward(UnfoldParams?params,?const?T*?in,?T*?out)?{
??CUDA_1D_KERNEL_LOOP_T(INDEX_T,?out_offset,?params.out_elem_cnt)?{
????using?ParamType?=?UnfoldParams;
????INDEX_T?in_index[ParamType::kInputNDim]?=?{0};
????INDEX_T?out_index[ParamType::kOutputNDim]?=?{0};
????params.out_index_helper.OffsetToNdIndex(out_offset,?out_index);
????if?(!UnfoldIndexTransform(params,?out_index,?in_index))?{
??????INDEX_T?in_offset?=?params.in_index_helper.NdIndexToOffset(in_index);
??????out[out_offset]?=?in[in_offset];
????}?else?{
??????out[out_offset]?=?static_cast(kUnfoldPaddingValue);
????}
??}
}

• 首先根據(jù)offset來計算當(dāng)前處理輸出元素的NdIndex
• 然后判斷UnfoldIndexTransform該方法的返回值
• 如果為false，則說明我們可以取到輸入元素，將其index轉(zhuǎn)換為1d的offset，并賦值給輸出
• 如果為true，則越界，我們填充先前設(shè)定好的一個padding_value(0)

至此整個img2col流程已經(jīng)完成，整體操作可參考下圖：

OneFlow版本的Fold

Fold就是將每一列填充回到kxk的地方

如果能理解前面的Unfold，那么這里的Fold也能很容易的理解。只是操作的元素反過來而已。

template<typename?T,?typename?INDEX_T,?int?NDIM,?int?SDIM>
__global__?void?CudaFoldForward(FoldParams?params,?const?T*?input_ptr,
????????????????????????????????T*?output_ptr)?{
??CUDA_1D_KERNEL_LOOP_T(INDEX_T,?in_offset,?params.in_elem_cnt)?{
????using?ParamType?=?FoldParams;
????INDEX_T?in_index[ParamType::kInputNDim]?=?{0};
????INDEX_T?out_index[ParamType::kOutputNDim]?=?{0};
????params.in_index_helper.OffsetToNdIndex(in_offset,?in_index);
????if?(!FoldIndexTransform(params,?in_index,?out_index))?{
??????INDEX_T?out_offset?=?params.out_index_helper.NdIndexToOffset(out_index);
??????XPUAdd::Invoke(&input_ptr[in_offset],?&output_ptr[out_offset]);
????}?else?{
??????continue;
????}
??}
}

沿用前面的索引映射邏輯，我們進入一個循環(huán)次數(shù)為輸入元素個數(shù)的循環(huán)體，計算得到當(dāng)前offset對應(yīng)的輸入NdIndex。

如果FoldIndexTransform返回的是false，則計算輸出的offset，并使用原子加atomic add，把輸入元素累加到該輸出位置。

小結(jié)

這部分代碼是接手同事寫一半的代碼完成的，不得不說同事的構(gòu)思真的很巧妙，通過模板參數(shù)能夠拓展1d 2d 3d，nchw, nhwc各種格式，盡管直觀上不太好理解。

darknet版本(Caffe同款)是比較直觀的幫助新手入門的img2col算法，可以結(jié)合那兩篇csdn博客來理解整個過程。