本文轉(zhuǎn)自徐飛翔的“緊致卷積網(wǎng)絡(luò)設(shè)計——Shift卷積算子”

版權(quán)聲明：本文為博主原創(chuàng)文章，遵循 CC 4.0 BY-SA 版權(quán)協(xié)議，轉(zhuǎn)載請附上原文出處鏈接和本聲明。

卷積計算及其優(yōu)化

為了討論的連續(xù)性，我們先簡單回顧下傳統(tǒng)的深度學習卷積計算。給定一個輸入張量，如Fig 1.1中的藍色塊所示，其尺寸為 ；給定卷積核 ，如Fig 1.1中的藍色虛線框所示，為了方便起見，假定步進stride = 1，padding = 1，那么最終得到輸出結(jié)果為 ，計算過程如式子(1.1)所示：

其中為卷積中心，而是卷積計算半徑的索引。不難知道，該卷積操作的參數(shù)量為。計算量也容易計算，考慮到每個卷積操作都需要對每個卷積核中的參數(shù)進行乘法計算，那么有乘法因子 ，而考慮到stride = 1而且存在填充，那么容易知道計算量為 ? FLOPs。容易發(fā)現(xiàn)，卷積的計算量和參數(shù)量與卷積核大小 ?呈現(xiàn)著二次增長的關(guān)系，這使得卷積的計算量和參數(shù)量增長都隨著網(wǎng)絡(luò)設(shè)計的加深變得難以控制。

Fig 1.1 經(jīng)典的卷積操作示意圖。

在進一步對傳統(tǒng)卷積計算進行優(yōu)化之前，我們先分析一下卷積計算到底提取了什么類型的信息。以二維卷積為例子，卷積計算主要在兩個維度提取信息，空間域和通道域，不過從本質(zhì)上說，通道域的信息可以看成是原始輸入（比如RGB圖片輸入）的層次化特征/信息的層疊，因此本質(zhì)上二維卷積還是提取空間域信息，只不過在層疊卷積過程中，使得空間域信息按照層次的特點，散布在了通道域中。

知道了這一點，我們就可以把卷積過程中的空間卷積和通道卷積分離開了，從而得到了所謂的 通道可分離卷積[4,5]。如Fig 1.2所示，這類型的卷積將空間域和通道域卷積完全分開，在第一步只考慮空間域卷積，因此對于每個輸入張量的通道，只會有唯一一個對應(yīng)的卷積核進行卷積。數(shù)學表達為：

對比式子(1.1)和(1.2)，我們發(fā)現(xiàn)區(qū)別在于對卷積核的索引上，通過式子(1.2)輸出的張量形狀為 ，為了接下來在通道域進行卷積，需要進一步應(yīng)用1x1卷積，將通道數(shù)從變?yōu)?，如式子(1.3)所示。

其中 為1x1卷積核。通道可分離卷積就是將傳統(tǒng)卷積(1.1)分解為了(1.2)(1.3)兩個步驟。

通過這種優(yōu)化，可以知道卷積核參數(shù)量變?yōu)?span>，而計算量變?yōu)?span>? FLOPs。雖然理論上，深度可分離網(wǎng)絡(luò)的確減少了計算量和參數(shù)量，但是實際上，因為深度可分離網(wǎng)絡(luò)的實現(xiàn)使得訪存1（memory access）過程占據(jù)了主導，使得實際計算占用率過小，限制了硬件的并行計算能力。

Fig 1.2 深度可分離卷積，對于輸入張量的每一個通道，都有其專有的卷積核進行卷積，最后通過一個1x1的卷積即可完成通道數(shù)量的放縮。

我們用傳統(tǒng)卷積和深度可分離卷積的計算/訪存系數(shù)進行比較（僅考慮最基本的訪存，即是將每個操作數(shù)都從內(nèi)存中獲取，而不考慮由于局部性原理[6]，而對重復的操作數(shù)進行訪問導致的消耗）：

式子(1.4)和(1.5)的比較最終會化簡為比較 和的大小，越小意味著計算效率越高。我們發(fā)現(xiàn)，傳統(tǒng)的卷積反而比深度可分離卷積的計算效率高得多。這是不利于程序并行計算的。

為此，文章[1]提出了Shift卷積算子，嘗試解決這種問題。Shift卷積算子

在Shift卷積算子中，其基本思路也是類似于深度可分離卷積的設(shè)計，將卷積分為空間域和通道域的卷積，通道域的卷積同樣是通過1x1卷積實現(xiàn)的，而在空間域卷積中，引入了shift操作。我們接下來會詳細地探討shift操作的設(shè)計啟發(fā)，細節(jié)和推導。

Fig 2.1 基于Shift的卷積可以分為Shift卷積算子和1x1卷積操作。

shift卷積算子的數(shù)學形式表達如式子(2.1)所示，如圖Fig 2.1所示，shift卷積的每一個卷積核都是一個“獨熱”的算子，其卷積核只有一個元素為1，其他全部為0，如式子(2.2)所示。類似于深度可分離卷積，對于輸入的 M個通道的張量，分別對應(yīng)了 M個Shift卷積核，如Fig 2.1的不同顏色的卷積核所示。

我們把其中一個通道的shift卷積操作拿出來分析，如Fig 2.2所示。我們發(fā)現(xiàn)，shift卷積過程相當于將原輸入的矩陣在某個方向進行平移，這也是為什么該操作稱之為shift的原因。雖然簡單的平移操作似乎沒有提取到空間信息，但是考慮到我們之前說到的，通道域是空間域信息的層次化擴散。因此通過設(shè)置不同方向的shift卷積核，可以將輸入張量不同通道進行平移，隨后配合1x1卷積實現(xiàn)跨通道的信息融合，即可實現(xiàn)空間域和通道域的信息提取。

Fig 2.2 shift卷積算子中的卷積操作，經(jīng)過填充后如（2）所示，我們發(fā)現(xiàn)，shift卷積相當于將原輸入矩陣在某個方向進行平移。

我們發(fā)現(xiàn)shift卷積的本質(zhì)是特定內(nèi)存的訪問，可學習參數(shù)只是集中在1x1卷積操作中。因此如果實現(xiàn)得當，shift卷積是不占用額外的計算量和參數(shù)量的，結(jié)合shift卷積，只使用1x1卷積即可提取到結(jié)構(gòu)化層次化的空間域信息，因此大大減少了卷積網(wǎng)絡(luò)設(shè)計的參數(shù)量和計算量。

然而我們注意到，對于一個卷積核大小為，通道數(shù)為M的卷積核而言，其可能的搜索空間為，在學習過程中窮盡這個搜索空間是不太現(xiàn)實的。為了減少搜索空間，[1]采用了一種簡單的啟發(fā)式設(shè)計：將 M個通道均勻地分成 個組，我們將每個組稱之為 平移組（shift group）。每個組有 個通道，這些通道都采用相同的平移方向。當然，有可能存在除不盡的情況，這個時候?qū)幸恍┩ǖ啦荒鼙粍澐值饺我庖粋€組內(nèi)，這些剩下的通道都稱之為“居中”組，如Fig 2.3所示，其中心元素為1，其他為0，也即是對原輸入不進行任何處理。

Fig 2.3 居中組的中心元素為1，其他元素都為0。

雖然通過這種手段大大縮小了搜索空間，但是仍然需要讓模型學出如何將第 m個通道映射到第個平移組的最佳排列規(guī)則，這仍然是一個很大的搜索空間。為了解決這個問題，以下需要提出一種方法，其能夠使得shift卷積層的輸出和輸入是關(guān)于通道排序無關(guān)的。假設(shè) 表示是在以為通道排序的shift卷積操作，那么公式(2.1)可以表示為 ，如果我們在進行該卷積之前，先后進行兩次通道排序，分別是 ??和 ，那么我們有：

其中表示算子組合。令 和分別表示1x1卷積操作，我們有式子(2.4)

這一點不難理解，即便對1x1卷積的輸入進行通道排序重組，在學習過程中，通過算法去調(diào)整1x1卷積的參數(shù)的順序，就可以通過構(gòu)造的方式，實現(xiàn) 和之間的雙射（bijective）。如式子(2.5)所示，就結(jié)論而言，不需要考慮通道的排序，比如只需要依次按著順序賦值某個平移組，使得其不重復即可。通過用1x1卷積“三明治”夾著shift卷積的操作，從理論上可以等價于其他任何形式的通道排序后的結(jié)果。這點比較繞，有疑問的讀者請在評論區(qū)留言。

根據(jù)以上討論，根據(jù)shift算子構(gòu)建出來的卷積模塊類似于Fig 2.4所示，注意到藍色實線塊的 1x1 conv -> shift kernel -> 1x1 conv正是和我們的討論一樣的結(jié)構(gòu)，而Identity塊則是考慮到仿照ResNet的設(shè)計補充的short cut鏈路。藍色虛線塊的shift塊是實驗補充的一個設(shè)計，存在虛線部分的shift塊的設(shè)計稱之為結(jié)構(gòu)，只存在實線部分的設(shè)計則稱之為 結(jié)構(gòu)。

Fig 2.4 基于shift卷積算子構(gòu)建的ResNet網(wǎng)絡(luò)基本模塊。

shift卷積算子的有效性在文章[1]設(shè)置了很多實驗進行對比，這里只給出證實其在分類任務(wù)上精度和計算量/參數(shù)量的一個比較，如Fig 2.5所示，我們發(fā)現(xiàn)shift算子的確在計算量和參數(shù)量上有著比較大的優(yōu)勢。

exp_resultFig 2.5 shift卷積網(wǎng)絡(luò)在CIFAR10/100分類任務(wù)上的表現(xiàn)對比表。

在[7]中有shift卷積算子前向和反向計算的cuda代碼，其主要操作就是進行卷積輸入張量的訪存選擇。有興趣的讀者可以自行移步去閱讀。

那么我只需要固定某個特定的索引順序即可，最簡單的方式如Fig a1所示，按行列排列的順序遍歷設(shè)置即可，并不需要對其進行shuffle，因為可以證實其本身都是可以通過結(jié)合1x1卷積的方式學習出來的（不同的索引順序?qū)W習出來的卷積參數(shù)不同，但是如果看成整體的話，它們是等價的）。因此文章里面應(yīng)該是不需要進行shift組的排序shuffle的。

shift_groupFig a1. 按順序排列的shift算子組示例。

說到如何實現(xiàn)shift的訪存優(yōu)化機制，我們可以先看看shift-gcn是怎么做實現(xiàn)的，具體見文章[8]。當然，本文并不是采用shift-gcn定義的那種shift圖卷積，我們回到開源的[7]中的具體代碼段進行分析。我不是很熟悉cuda編程，只能作初步的分析，比如代碼段[9]：

__global__ void shiftnet_cuda_moduloshift3x3_nchw_float32_kernel_tilein16x16_tileout14x14(
    float *src,
    float *dst,
    int num_h_tiles,
    int num_w_tiles,
    int batch_sz,
    int channels,
    int height,
    int width)
{
  __shared__ float cache[256];
  const int num_blocks = batch_sz * channels * num_h_tiles * num_w_tiles;
  const int num_threads = blockDim.x * num_blocks;
  const int rd_chans = (channels / 9) * 9;
  for (int idx = threadIdx.x + blockDim.x * blockIdx.x;
      idx < num_threads; idx += blockDim.x * gridDim.x)
  {
    const int w_tile_idx = (idx / 256) % num_w_tiles;
    const int h_tile_idx = ((idx / 256) / num_w_tiles) % num_h_tiles;
    const int tile_ch = (((idx / 256) / num_w_tiles) / num_h_tiles) % channels;
    const int tile_batch_idx = ((((idx / 256) / num_w_tiles) / num_h_tiles) / channels) % batch_sz;
    const int w_shift = ((tile_ch % 3) - 1) * (tile_ch < rd_chans);
    const int h_shift = (((tile_ch / 3) % 3) - 1) * (tile_ch < rd_chans);
    const int w_tile_off = threadIdx.x % 16;
    const int h_tile_off = threadIdx.x / 16;
    const int w_idx = w_tile_off - 1 + 14 * w_tile_idx;
    const int h_idx = h_tile_off - 1 + 14 * h_tile_idx;
    const int buf_idx = w_idx + width * (h_idx + height * (tile_ch + channels * tile_batch_idx));
    if (w_idx >= 0 && w_idx < width && h_idx >= 0 && h_idx < height) {
      cache[threadIdx.x] = src[buf_idx];
    } else {
      cache[threadIdx.x] = 0.0f;
    }
    __syncthreads();
    if (w_tile_off >= 1 && w_tile_off < 15 && h_tile_off >= 1 && h_tile_off < 15) {
      if (w_idx >= 0 && w_idx < width && h_idx >= 0 && h_idx < height) {
        const int cache_idx = (w_tile_off + w_shift) + 16 * (h_tile_off + h_shift);
        dst[buf_idx] = cache[cache_idx];
      }
    }
    __syncthreads();
  }
}

它的實現(xiàn)并沒有完全優(yōu)化完全，因為他沒有結(jié)合后續(xù)的1x1卷積進行優(yōu)化，他只是進行了將某行（或列）置位0，代碼是:

if (w_idx >= 0 && w_idx < width && h_idx >= 0 && h_idx < height) {
      cache[threadIdx.x] = src[buf_idx];
    } else {
      cache[threadIdx.x] = 0.0f;
    }
   }

然后進行移位：

    if (w_tile_off >= 1 && w_tile_off < 15 && h_tile_off >= 1 && h_tile_off < 15) {
      if (w_idx >= 0 && w_idx < width && h_idx >= 0 && h_idx < height) {
        const int cache_idx = (w_tile_off + w_shift) + 16 * (h_tile_off + h_shift);
        dst[buf_idx] = cache[cache_idx];
      }
    }

我覺得這個并不是最優(yōu)化后的結(jié)果，最優(yōu)化應(yīng)該是指定某個方塊（比如不包括第一列的其他所有數(shù)據(jù)），與后續(xù)的1x1卷積聯(lián)合起來，只對這些方塊進行卷積，這才是真正的訪存優(yōu)化，顯然這樣難度太大，因此它的實現(xiàn)并沒有這樣做。（正如我所說的，我不是很熟悉cuda，有謬誤請指出。）

最后需要指出的是，它并不是one-hot矩陣點乘，而是卷積。

以上。

Reference

[1]. Wu, B., Wan, A., Yue, X., Jin, P., Zhao, S., Golmant, N., … & Keutzer, K. (2018). Shift: A zero flop, zero parameter alternative to spatial convolutions. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (pp. 9127-9135).

[2]. Cheng, K., Zhang, Y., He, X., Chen, W., Cheng, J., & Lu, H. (2020). Skeleton-Based Action Recognition With Shift Graph Convolutional Network. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 183-192).

[3]. https://github.com/peterhj/shiftnet_cuda_v2

[4]. Howard, Andrew G., Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. “Mobilenets: Efficient convolutional neural networks for mobile vision applications.” arXiv preprint arXiv:1704.04861 (2017).

[5]. Chollet, F. (2017). Xception: Deep learning with depthwise separable convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 1251-1258).

[6]. https://baike.baidu.com/item/%E5%B1%80%E9%83%A8%E6%80%A7%E5%8E%9F%E7%90%86

[7]. https://github.com/peterhj/shiftnet_cuda_v2/blob/master/src/shiftnet_cuda_kernels.cu

[8]. https://fesian.blog.csdn.net/article/details/109644297

[9]. https://github.com/peterhj/shiftnet_cuda_v2/blob/4d471bd744751ff0fd6cf5acd518e9484cc70a98/src/shiftnet_cuda_kernels.cu#L25