紧致卷积网络设计——Shift卷积算子

FesianXu 2020/10/29 at UESTC

前言

最近笔者在阅读关于骨骼点数据动作识别的文献Shift-GCN[2]的时候，发现了原来还有Shift卷积算子[1]这种东西，该算子是一种可供作为空间卷积的替代品，其理论上不需要增添额外的计算量和参数量，就可以通过1x1卷积实现空间域和通道域的卷积，是一种做紧致模型设计的好工具。本文作为笔记纪录笔者的论文阅读思考， 如有谬误请联系指出，转载请联系作者并注明出处，谢谢。

$\nabla$ 联系方式：

e-mail: FesianXu@gmail.com

QQ: 973926198

github: github.com/FesianXu

知乎专栏: 计算机视觉/计算机图形理论与应用

微信公众号： qrcode

卷积计算及其优化

为了讨论的连续性，我们先简单回顾下传统的深度学习卷积计算。给定一个输入张量，如Fig 1.1中的蓝色块所示，其尺寸为 $\mathbf{F} \in \mathbb{R}^{D_F \times D_F \times M}$ ；给定卷积核 $\mathbf{K} \in \mathbb{R}^{D_K \times D_K \times M \times N}$ ，如Fig 1.1中的蓝色虚线框所示，为了方便起见，假定步进stride = 1，padding = 1，那么最终得到输出结果为 $\mathbf{G} \in \mathbb{R}^{D_F \times D_F \times N}$ ，计算过程如式子(1.1)所示：

G_{k,l,n} = \sum_{i,j,m} K_{i,j,m,n}F_{k+\hat{i},l+\hat{j},m} \tag{1.1}

其中 $(k,l)$ 为卷积中心，而 $\hat{i} = i-\lfloor D_K/2\rfloor$ ， $\hat{j} = j-\lfloor D_K /2 \rfloor$ 是卷积计算半径的索引。不难知道，该卷积操作的参数量为 $M \times N \times D_K^2$ 。计算量也容易计算，考虑到每个卷积操作都需要对每个卷积核中的参数进行乘法计算，那么有乘法因子 $D_K^2$ ，而考虑到stride = 1而且存在填充，那么容易知道计算量为 $M \times N \times D_F^2 \times D_K^2$ FLOPs。容易发现，卷积的计算量和参数量与卷积核大小 $D_K$ 呈现着二次增长的关系，这使得卷积的计算量和参数量增长都随着网络设计的加深变得难以控制。

convolution_operator

Fig 1.1 经典的卷积操作示意图。

在进一步对传统卷积计算进行优化之前，我们先分析一下卷积计算到底提取了什么类型的信息。以二维卷积为例子，卷积计算主要在两个维度提取信息，空间域和通道域，不过从本质上说，通道域的信息可以看成是原始输入（比如RGB图片输入）的层次化特征/信息的层叠，因此本质上二维卷积还是提取空间域信息，只不过在层叠卷积过程中，使得空间域信息按照层次的特点，散布在了通道域中。

知道了这一点，我们就可以把卷积过程中的空间卷积和通道卷积分离开了，从而得到了所谓的 通道可分离卷积[4,5]。如Fig 1.2所示，这类型的卷积将空间域和通道域卷积完全分开，在第一步只考虑空间域卷积，因此对于每个输入张量的通道，只会有唯一一个对应的卷积核进行卷积。数学表达为：

\hat{G}_{k,l,m} = \sum_{i,j} \hat{K}_{i,j,m} F_{k+\hat{i},l+\hat{j},m} \tag{1.2}

对比式子(1.1)和(1.2)，我们发现区别在于对卷积核的索引上，通过式子(1.2)输出的张量形状为 $\hat{\mathbf{G}} \in \mathbb{R}^{D_F \times D_F \times M}$ ，为了接下来在通道域进行卷积，需要进一步应用1x1卷积，将通道数从 $M$ 变为 $N$ ，如式子(1.3)所示。

G_{k,l,n} = \sum_{m} P_{m,n} \hat{G}_{k,l,m} \tag{1.3}

其中 $P \in \mathbb{R}^{M \times N}$ 为1x1卷积核。通道可分离卷积就是将传统卷积(1.1)分解为了(1.2)(1.3)两个步骤。

通过这种优化，可以知道卷积核参数量变为 $M \times D_K^2$ ，而计算量变为 $M \times D_K^2 \times D_F^2$ FLOPs。虽然理论上，深度可分离网络的确减少了计算量和参数量，但是实际上，因为深度可分离网络的实现使得访存¹（memory access）过程占据了主导，使得实际计算占用率过小，限制了硬件的并行计算能力。

depth_wise_conv

Fig 1.2 深度可分离卷积，对于输入张量的每一个通道，都有其专有的卷积核进行卷积，最后通过一个1x1的卷积即可完成通道数量的放缩。

我们用传统卷积和深度可分离卷积的计算/访存系数进行比较（仅考虑最基本的访存，即是将每个操作数都从内存中获取，而不考虑由于局部性原理[6]，而对重复的操作数进行访问导致的消耗）：

传统卷积:\dfrac{M \times N \times D_F^2 \times D_K^2}{D_F^2 \times (M+N) + D_F^2 \times M \times N} \tag{1.4}

深度可分离卷积: \dfrac{M \times D_F^2 \times D_K^2}{D_F^2 \times 2M + D_K^2 \times M} \tag{1.5}

式子(1.4)和(1.5)的比较最终会化简为比较 $(M+N)/N$ 和 $2M$ 的大小，越小意味着计算效率越高。我们发现，传统的卷积反而比深度可分离卷积的计算效率高得多。这是不利于程序并行计算的。

为此，文章[1]提出了Shift卷积算子，尝试解决这种问题。

Shift卷积算子

在Shift卷积算子中，其基本思路也是类似于深度可分离卷积的设计，将卷积分为空间域和通道域的卷积，通道域的卷积同样是通过1x1卷积实现的，而在空间域卷积中，引入了shift操作。我们接下来会详细地探讨shift操作的设计启发，细节和推导。

shift_conv

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

shift卷积算子的数学形式表达如式子(2.1)所示，如图Fig 2.1所示，shift卷积的每一个卷积核都是一个“独热”的算子，其卷积核只有一个元素为1，其他全部为0，如式子(2.2)所示。类似于深度可分离卷积，对于输入的 $M$ 个通道的张量，分别对应了 $M$ 个Shift卷积核，如Fig 2.1的不同颜色的卷积核所示。

\tilde{G}_{k,l,m} = \sum_{i,j} \tilde{K}_{i,j,m} F_{k+\hat{i},l+\hat{j},m} \tag{2.1}

\tilde{K}_{i,j,m} = \left\{ \begin{aligned} 1, & & 当 i = i_m，j=j_m \\ 0, & & 其他 \end{aligned} \right. \tag{2.2}

我们把其中一个通道的shift卷积操作拿出来分析，如Fig 2.2所示。我们发现，shift卷积过程相当于将原输入的矩阵在某个方向进行平移，这也是为什么该操作称之为shift的原因。虽然简单的平移操作似乎没有提取到空间信息，但是考虑到我们之前说到的，通道域是空间域信息的层次化扩散。因此通过设置不同方向的shift卷积核，可以将输入张量不同通道进行平移，随后配合1x1卷积实现跨通道的信息融合，即可实现空间域和通道域的信息提取。

shift_conv_detail

Fig 2.2 shift卷积算子中的卷积操作，经过填充后如（2）所示，我们发现，shift卷积相当于将原输入矩阵在某个方向进行平移。

我们发现shift卷积的本质是特定内存的访问，可学习参数只是集中在1x1卷积操作中。因此如果实现得当，shift卷积是不占用额外的计算量和参数量的，结合shift卷积，只使用1x1卷积即可提取到结构化层次化的空间域信息，因此大大减少了卷积网络设计的参数量和计算量。

然而我们注意到，对于一个卷积核大小为 $D_K$ ，通道数为 $M$ 的卷积核而言，其可能的搜索空间为 $(D_K^2)^M$ ，在学习过程中穷尽这个搜索空间是不太现实的。为了减少搜索空间，[1]采用了一种简单的启发式设计：将 $M$ 个通道均匀地分成 $D_K^2$ 个组，我们将每个组称之为 平移组（shift group）。每个组有 $\lfloor M/D_K^2\rfloor$ 个通道，这些通道都采用相同的平移方向。当然，有可能存在除不尽的情况，这个时候将会有一些通道不能被划分到任意一个组内，这些剩下的通道都称之为“居中”组，如Fig 2.3所示，其中心元素为1，其他为0，也即是对原输入不进行任何处理。

center_group

Fig 2.3 居中组的中心元素为1，其他元素都为0。

虽然通过这种手段大大缩小了搜索空间，但是仍然需要让模型学出如何将第 $m$ 个通道映射到第 $n, n\in[0,\lfloor M/D_K^2\rfloor-1]$ 个平移组的最佳排列规则，这仍然是一个很大的搜索空间。为了解决这个问题，以下需要提出一种方法，其能够使得shift卷积层的输出和输入是关于通道排序无关的。假设 $\mathcal{K}_{\pi}(\cdot)$ 表示是在以 $\pi$ 为通道排序的shift卷积操作，那么公式(2.1)可以表示为 $\tilde{G} = \mathcal{K}_{\pi}(F)$ ，如果我们在进行该卷积之前，先后进行两次通道排序，分别是 $\mathcal{P}_{\pi_1}$ 和 $\mathcal{P}_{\pi_2}$ ，那么我们有：

\tilde{G} = \mathcal{P}_{\pi_2}(\mathcal{K}_{\pi}(\mathcal{P}_{\pi_1}(F))) = (\mathcal{P}_{\pi_2} \circ \mathcal{K}_{\pi} \circ \mathcal{P}_{\pi_1})(F) \tag{2.3}

其中 $\circ$ 表示算子组合。令 $\mathcal{P}_1(\cdot)$ 和 $\mathcal{P}_2(\cdot)$ 分别表示1x1卷积操作，我们有式子(2.4)

\begin{aligned} \hat{P}_1 &= \mathcal{P}_1 \circ \mathcal{P}_{\pi_1} \\ \hat{P}_2 &= \mathcal{P}_2 \circ \mathcal{P}_{\pi_2} \end{aligned} \tag{2.4}

这一点不难理解，即便对1x1卷积的输入进行通道排序重组，在学习过程中，通过算法去调整1x1卷积的参数的顺序，就可以通过构造的方式，实现 $\hat{\mathcal{P}}_{x}$ 和 $\mathcal{P}_{x}$ 之间的双射（bijective）。如式子(2.5)所示，就结论而言，不需要考虑通道的排序，比如只需要依次按着顺序赋值某个平移组，使得其不重复即可。通过用1x1卷积“三明治”夹着shift卷积的操作，从理论上可以等价于其他任何形式的通道排序后的结果。这点比较绕，有疑问的读者请在评论区留言。

\begin{aligned} G &= (\mathcal{P}_{2} \circ \mathcal{P}_{\pi_2} \circ \mathcal{K}_{\pi} \circ \mathcal{P}_{\pi_1} \circ \mathcal{P}_1)(F) \\ &= ((\mathcal{P}_{2} \circ \mathcal{P}_{\pi_2}) \circ \mathcal{K}_{\pi} \circ (\mathcal{P}_{\pi_1} \circ \mathcal{P}_1))(F) \\ &= (\hat{\mathcal{P}}_2 \circ \mathcal{K}_{\pi} \circ \hat{\mathcal{P}}_1)(F) \end{aligned} \tag{2.5}

根据以上讨论，根据shift算子构建出来的卷积模块类似于Fig 2.4所示，注意到蓝色实线块的1x1 conv -> shift kernel -> 1x1 conv正是和我们的讨论一样的结构，而Identity块则是考虑到仿照ResNet的设计补充的short cut链路。蓝色虚线块的shift块是实验补充的一个设计，存在虚线部分的shift块的设计称之为 $SC^2$ 结构，只存在实线部分的设计则称之为 $CSC$ 结构。

shift_resnet_block

Fig 2.4 基于shift卷积算子构建的ResNet网络基本模块。

shift卷积算子的有效性在文章[1]设置了很多实验进行对比，这里只给出证实其在分类任务上精度和计算量/参数量的一个比较，如Fig 2.5所示，我们发现shift算子的确在计算量和参数量上有着比较大的优势。

exp_result

Fig 2.5 shift卷积网络在CIFAR10/100分类任务上的表现对比表。

在[7]中有shift卷积算子前向和反向计算的cuda代码，其主要操作就是进行卷积输入张量的访存选择。有兴趣的读者可以自行移步去阅读。

Update 20201115：在这里回答一位来自于知乎朋友的问题，链接来源：zhuanlan.zhihu.com/p/272732283

ID:fighting shift操作理论上就是点乘一个one-hot矩阵（只有一个位置为1，其余位置为0），M个channel的shift kernel为D_K^2^M，为了降低搜索空间，对M个channel分组为M/D_K^2，每个组内采用一个one-shot矩阵。也即是说仅需要将M个channel中的每一个channel给出它属于这些组中的一个组。为了使得这么多group channel在输入输出时排序无关（仅对每一个channel考虑将其映射为哪一个shift 矩阵，但是与这个shift矩阵所在的group channel的位置没有关系，也就是所有的group channel地位一样，仅仅考虑一种排序方式即可，比如按照划窗的顺序从上到下的group channel对应的shift矩阵），就人为的shuffle channel，但shuffle channel不可导，将shuffle操作和11卷积相结合。所以最终推导出来的式子含义：先对输入进行11卷积和channel shuffle，然后执行某种特定排序的shift 移动，再对输出执行shuffle channel和11卷积。两个疑问：1，对输出还要在执行shuffle和11卷积，是为了啥呢？2，shift 操作是操作内存实现的，怎么实现的呢？

$\rightarrow$ 回答：你好，我认为这里其实并不需要对通道进行shuffle打乱，因为既然根据式子(2.3-2.5)的推导可以得知

这么多group channel在输入输出时排序无关

那么我只需要固定某个特定的索引顺序即可，最简单的方式如Fig a1所示，按行列排列的顺序遍历设置即可，并不需要对其进行shuffle，因为可以证实其本身都是可以通过结合1x1卷积的方式学习出来的（不同的索引顺序学习出来的卷积参数不同，但是如果看成整体的话，它们是等价的）。因此文章里面应该是不需要进行shift组的排序shuffle的。

shift_group

Fig 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]. github.com/peterhj/shi…

[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]. baike.baidu.com/item/%E5%B1…

[7]. github.com/peterhj/shi…

[8]. fesian.blog.csdn.net/article/det…

[9]. github.com/peterhj/shi…

访存指的是从内存中取出操作数加载到寄存器中，通常访存时间远比计算时间长，大概是数量级上的差别。 ↩

紧致卷积网络设计——Shift卷积算子

前言

卷积计算及其优化

Shift卷积算子

Reference

Footnotes