本文轉自徐飛翔的“緊致卷積網絡設計——Shift卷積算子”

版權聲明：本文為博主原創文章，遵循 CC 4.0 BY-SA 版權協議，轉載請附上原文出處鏈接和本聲明。

卷積計算及其優化

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

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

Fig 1.1 經典的卷積操作示意圖。

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

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

對比式子(1.1)和(1.2)，我們發現區別在于對卷積核的索引上，通過式子(1.2)輸出的張量形狀為 ，為了接下來在通道域進行卷積，需要進一步應用1x1卷積，將通道數從變為 ，如式子(1.3)所示。

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

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

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

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

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

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

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

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

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

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

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

我們發現shift卷積的本質是特定內存的訪問，可學習參數只是集中在1x1卷積操作中。因此如果實現得當，shift卷積是不占用額外的計算量和參數量的，結合shift卷積，只使用1x1卷積即可提取到結構化層次化的空間域信息，因此大大減少了卷積網絡設計的參數量和計算量。

然而我們注意到，對于一個卷積核大小為，通道數為M的卷積核而言，其可能的搜索空間為，在學習過程中窮盡這個搜索空間是不太現實的。為了減少搜索空間，[1]采用了一種簡單的啟發式設計：將 M個通道均勻地分成 個組，我們將每個組稱之為 平移組（shift group）。每個組有 個通道，這些通道都采用相同的平移方向。當然，有可能存在除不盡的情況，這個時候將會有一些通道不能被劃分到任意一個組內，這些剩下的通道都稱之為“居中”組，如Fig 2.3所示，其中心元素為1，其他為0，也即是對原輸入不進行任何處理。

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

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

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

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

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

Fig 2.4 基于shift卷積算子構建的ResNet網絡基本模塊。

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

exp_resultFig 2.5 shift卷積網絡在CIFAR10/100分類任務上的表現對比表。

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

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

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

說到如何實現shift的訪存優化機制，我們可以先看看shift-gcn是怎么做實現的，具體見文章[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();
  }
}

它的實現并沒有完全優化完全，因為他沒有結合后續的1x1卷積進行優化，他只是進行了將某行（或列）置位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];
      }
    }

我覺得這個并不是最優化后的結果，最優化應該是指定某個方塊（比如不包括第一列的其他所有數據），與后續的1x1卷積聯合起來，只對這些方塊進行卷積，這才是真正的訪存優化，顯然這樣難度太大，因此它的實現并沒有這樣做。（正如我所說的，我不是很熟悉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