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

紧致卷积网络设计——Shift卷积算子
FesianXu 2020/10/29 at UESTC

前言

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

∇ \nabla 联系方式:

e-mail: [email protected]

QQ: 973926198

github: https://github.com/FesianXu

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


卷积计算及其优化

为了讨论的连续性,我们先简单回顾下传统的深度学习卷积计算。给定一个输入张量,如Fig 1.1中的蓝色块所示,其尺寸为 F ∈ R D F × D F × M \mathbf{F} \in \mathbb{R}^{D_F \times D_F \times M} FRDF×DF×M;给定卷积核 K ∈ R D K × D K × M × N \mathbf{K} \in \mathbb{R}^{D_K \times D_K \times M \times N} KRDK×DK×M×N,如Fig 1.1中的蓝色虚线框所示,为了方便起见,假定步进stride = 1padding = 1,那么最终得到输出结果为 G ∈ R D F × D F × N \mathbf{G} \in \mathbb{R}^{D_F \times D_F \times N} GRDF×DF×N,计算过程如式子(1.1)所示:
G k , l , n = ∑ i , j , m K i , j , m , n F k + i ^ , l + j ^ , m (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} Gk,l,n=i,j,mKi,j,m,nFk+i^,l+j^,m(1.1)
其中 ( k , l ) (k,l) (k,l)为卷积中心,而 i ^ = i − ⌊ D K / 2 ⌋ \hat{i} = i-\lfloor D_K/2\rfloor i^=iDK/2 j ^ = j − ⌊ D K / 2 ⌋ \hat{j} = j-\lfloor D_K /2 \rfloor j^=jDK/2是卷积计算半径的索引。不难知道,该卷积操作的参数量为 M × N × D K 2 M \times N \times D_K^2 M×N×DK2。计算量也容易计算,考虑到每个卷积操作都需要对每个卷积核中的参数进行乘法计算,那么有乘法因子 D K 2 D_K^2 DK2,而考虑到stride = 1而且存在填充,那么容易知道计算量为 M × N × D F 2 × D K 2 M \times N \times D_F^2 \times D_K^2 M×N×DF2×DK2 FLOPs。容易发现,卷积的计算量和参数量与卷积核大小 D K D_K DK呈现着二次增长的关系,这使得卷积的计算量和参数量增长都随着网络设计的加深变得难以控制。

convolution_operator

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

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

知道了这一点,我们就可以把卷积过程中的空间卷积和通道卷积分离开了,从而得到了所谓的 通道可分离卷积[4,5]。如Fig 1.2所示,这类型的卷积将空间域和通道域卷积完全分开,在第一步只考虑空间域卷积,因此对于每个输入张量的通道,只会有唯一一个对应的卷积核进行卷积。数学表达为:
G ^ k , l , m = ∑ i , j K ^ i , j , m F k + i ^ , l + j ^ , m (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} G^k,l,m=i,jK^i,j,mFk+i^,l+j^,m(1.2)
对比式子(1.1)和(1.2),我们发现区别在于对卷积核的索引上,通过式子(1.2)输出的张量形状为 G ^ ∈ R D F × D F × M \hat{\mathbf{G}} \in \mathbb{R}^{D_F \times D_F \times M} G^RDF×DF×M,为了接下来在通道域进行卷积,需要进一步应用1x1卷积,将通道数从 M M M变为 N N N,如式子(1.3)所示。
G k , l , n = ∑ m P m , n G ^ k , l , m (1.3) G_{k,l,n} = \sum_{m} P_{m,n} \hat{G}_{k,l,m} \tag{1.3} Gk,l,n=mPm,nG^k,l,m(1.3)
其中 P ∈ R M × N P \in \mathbb{R}^{M \times N} PRM×N为1x1卷积核。通道可分离卷积就是将传统卷积(1.1)分解为了(1.2)(1.3)两个步骤。

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

depth_wise_conv

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

我们用传统卷积和深度可分离卷积的计算/访存系数进行比较(仅考虑最基本的访存,即是将每个操作数都从内存中获取,而不考虑由于局部性原理[6],而对重复的操作数进行访问导致的消耗):
传 统 卷 积 : M × N × D F 2 × D K 2 D F 2 × ( M + N ) + D F 2 × M × N (1.4) 传统卷积:\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} :DF2×(M+N)+DF2×M×NM×N×DF2×DK2(1.4)

深 度 可 分 离 卷 积 : M × D F 2 × D K 2 D F 2 × 2 M + D K 2 × M (1.5) 深度可分离卷积: \dfrac{M \times D_F^2 \times D_K^2}{D_F^2 \times 2M + D_K^2 \times M} \tag{1.5} :DF2×2M+DK2×MM×DF2×DK2(1.5)

式子(1.4)和(1.5)的比较最终会化简为比较 ( M + N ) / N (M+N)/N (M+N)/N 2 M 2M 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 M个通道的张量,分别对应了 M M M个Shift卷积核,如Fig 2.1的不同颜色的卷积核所示。

G ~ k , l , m = ∑ i , j K ~ i , j , m F k + i ^ , l + j ^ , m (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} G~k,l,m=i,jK~i,j,mFk+i^,l+j^,m(2.1)

K ~ i , j , m = { 1 , 当 i = i m , j = j m 0 , 其 他 (2.2) \tilde{K}_{i,j,m} = \left\{ \begin{aligned} 1, & & 当 i = i_m,j=j_m \\ 0, & & 其他 \end{aligned} \right. \tag{2.2} K~i,j,m={1,0,i=imj=jm(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 D_K DK,通道数为 M M M的卷积核而言,其可能的搜索空间为 ( D K 2 ) M (D_K^2)^M (DK2)M,在学习过程中穷尽这个搜索空间是不太现实的。为了减少搜索空间,[1]采用了一种简单的启发式设计:将 M M M个通道均匀地分成 D K 2 D_K^2 DK2个组,我们将每个组称之为 平移组(shift group)。每个组有 ⌊ M / D K 2 ⌋ \lfloor M/D_K^2\rfloor M/DK2个通道,这些通道都采用相同的平移方向。当然,有可能存在除不尽的情况,这个时候将会有一些通道不能被划分到任意一个组内,这些剩下的通道都称之为“居中”组,如Fig 2.3所示,其中心元素为1,其他为0,也即是对原输入不进行任何处理。

center_group

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

虽然通过这种手段大大缩小了搜索空间,但是仍然需要让模型学出如何将第 m m m个通道映射到第 n , n ∈ [ 0 , ⌊ M / D K 2 ⌋ − 1 ] n, n\in[0,\lfloor M/D_K^2\rfloor-1] n,n[0,M/DK21]个平移组的最佳排列规则,这仍然是一个很大的搜索空间。为了解决这个问题,以下需要提出一种方法,其能够使得shift卷积层的输出和输入是关于通道排序无关的。假设 K π ( ⋅ ) \mathcal{K}_{\pi}(\cdot) Kπ()表示是在以 π \pi π为通道排序的shift卷积操作,那么公式(2.1)可以表示为 G ~ = K π ( F ) \tilde{G} = \mathcal{K}_{\pi}(F) G~=Kπ(F),如果我们在进行该卷积之前,先后进行两次通道排序,分别是 P π 1 \mathcal{P}_{\pi_1} Pπ1 P π 2 \mathcal{P}_{\pi_2} Pπ2,那么我们有:
G ~ = P π 2 ( K π ( P π 1 ( F ) ) ) = ( P π 2 ∘ K π ∘ P π 1 ) ( F ) (2.3) \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} G~=Pπ2(Kπ(Pπ1(F)))=(Pπ2KπPπ1)(F)(2.3)
其中 ∘ \circ 表示算子组合。令 P 1 ( ⋅ ) \mathcal{P}_1(\cdot) P1() P 2 ( ⋅ ) \mathcal{P}_2(\cdot) P2()分别表示1x1卷积操作,我们有式子(2.4)
P ^ 1 = P 1 ∘ P π 1 P ^ 2 = P 2 ∘ P π 2 (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} P^1P^2=P1Pπ1=P2Pπ2(2.4)
这一点不难理解,即便对1x1卷积的输入进行通道排序重组,在学习过程中,通过算法去调整1x1卷积的参数的顺序,就可以通过构造的方式,实现 P ^ x \hat{\mathcal{P}}_{x} P^x P x \mathcal{P}_{x} Px之间的双射(bijective)。如式子(2.5)所示,就结论而言,不需要考虑通道的排序,比如只需要依次按着顺序赋值某个平移组,使得其不重复即可。通过用1x1卷积“三明治”夹着shift卷积的操作,从理论上可以等价于其他任何形式的通道排序后的结果。这点比较绕,有疑问的读者请在评论区留言。
G = ( P 2 ∘ P π 2 ∘ K π ∘ P π 1 ∘ P 1 ) ( F ) = ( ( P 2 ∘ P π 2 ) ∘ K π ∘ ( P π 1 ∘ P 1 ) ) ( F ) = ( P ^ 2 ∘ K π ∘ P ^ 1 ) ( F ) (2.5) \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} G=(P2Pπ2KπPπ1P1)(F)=((P2Pπ2)Kπ(Pπ1P1))(F)=(P^2KπP^1)(F)(2.5)
根据以上讨论,根据shift算子构建出来的卷积模块类似于Fig 2.4所示,注意到蓝色实线块的1x1 conv -> shift kernel -> 1x1 conv正是和我们的讨论一样的结构,而Identity块则是考虑到仿照ResNet的设计补充的short cut链路。蓝色虚线块的shift块是实验补充的一个设计,存在虚线部分的shift块的设计称之为 S C 2 SC^2 SC2结构,只存在实线部分的设计则称之为 C S C CSC 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代码,其主要操作就是进行卷积输入张量的访存选择。有兴趣的读者可以自行移步去阅读。


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


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