引言:一次小小的复习
卷积和池化的 img2col是怎么计算的?,反向传播过程又有哪些需要注意的地方? 维度变换到底是怎么一回事? 如果你感兴趣不妨看下去 -qwe
众所周知,由于传统的卷积计算速度太慢,所以有了img2col技术,那我这里假设读者已经都是对卷积和池化两个概念非常熟悉的同学了,我们本文就专注于这两个算子在代码上的实现,在实现之前也带大家复习一遍简单的推导。其中可能有一些错误的地方,求轻喷和大佬们的谅解!
前向计算的部分,比较基础,大家应该都会,我们这里给出前向计算的img2col版本代码,这部分代码是我的自制框架中的一部份,感兴趣的话可以看我的github。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 class Conv (TensorOp ): def __init__ (self, stride: Optional [int ] = 1 , padding: Optional [int ] = 0 ): self.stride = stride self.padding = padding def compute (self, A, B ): N, H, W, C_in = A.shape K, _, I, C_out = B.shape assert C_in == I, "input tensor shape and kernel doesn't match" pad_width = [(0 , 0 ), (self.padding, self.padding), (self.padding, self.padding), (0 , 0 )] _A = array_api.pad(A, pad_width, mode='constant' , constant_values=0 ) if self.padding > 0 else A inner_dim = K * K * C_in Ns, Hs, Ws, Cs = _A.strides H_out = (H - K + 2 * self.padding) // self.stride + 1 W_out = (W - K + 2 * self.padding) // self.stride + 1 _A = array_api.as_strided( _A, shape=(N, H_out, W_out, K, K, C_in), strides=(Ns, Hs*self.stride, Ws*self.stride, Hs, Ws, Cs) ) _A_ = compact(_A).reshape((-1 , inner_dim)) _B_ = compact(B).reshape((-1 , C_out)) out = _A_ @ _B_ return out.reshape((N, H_out, W_out, C_out))
现在我们重点关注卷积反向传播部分的代码,如何高效的实现卷积的反向传播。为了便于理解,我们从一个简单的例子开始。考虑如下的卷积操作:
[ X 00 X 01 X 02 X 10 X 11 X 12 X 20 X 21 X 22 ] ∗ [ W 00 W 01 W 10 W 11 ] = [ Z 00 Z 01 Z 10 Z 11 ] \begin{equation}
\begin{bmatrix}
X_{00} & X_{01} & X_{02} \\
X_{10} & X_{11} & X_{12} \\
X_{20} & X_{21} & X_{22}
\end{bmatrix}
\ast
\begin{bmatrix}
W_{00} & W_{01} \\
W_{10} & W_{11}
\end{bmatrix}
=
\begin{bmatrix}
Z_{00} & Z_{01} \\
Z_{10} & Z_{11}
\end{bmatrix}
\end{equation}
X 00 X 10 X 20 X 01 X 11 X 21 X 02 X 12 X 22 ∗ [ W 00 W 10 W 01 W 11 ] = [ Z 00 Z 10 Z 01 Z 11 ]
X 矩阵作为我们的输入,W 矩阵作为我们的卷积核(参数矩阵),Z 矩阵作为我们的输出矩阵。没有padding,
stride默认为1,输入输出通道数都是1。于是我们可以得到如下的计算公式。
{ Z 00 = X 00 W 00 + X 01 W 01 + X 10 W 10 + X 11 W 11 Z 01 = X 01 W 00 + X 02 W 01 + X 11 W 10 + X 12 W 11 Z 10 = X 10 W 00 + X 11 W 01 + X 20 W 10 + X 21 W 11 Z 11 = X 11 W 00 + X 12 W 01 + X 21 W 10 + X 22 W 11 \begin{equation}
\begin{cases}
Z_{00} = X_{00}W_{00} + X_{01}W_{01} + X_{10}W_{10} + X_{11}W_{11} \\
Z_{01} = X_{01}W_{00} + X_{02}W_{01} + X_{11}W_{10} + X_{12}W_{11} \\
Z_{10} = X_{10}W_{00} + X_{11}W_{01} + X_{20}W_{10} + X_{21}W_{11} \\
Z_{11} = X_{11}W_{00} + X_{12}W_{01} + X_{21}W_{10} + X_{22}W_{11} \\
\end{cases}
\end{equation}
⎩ ⎨ ⎧ Z 00 = X 00 W 00 + X 01 W 01 + X 10 W 10 + X 11 W 11 Z 01 = X 01 W 00 + X 02 W 01 + X 11 W 10 + X 12 W 11 Z 10 = X 10 W 00 + X 11 W 01 + X 20 W 10 + X 21 W 11 Z 11 = X 11 W 00 + X 12 W 01 + X 21 W 10 + X 22 W 11
现在我们思考对于这样的一个卷积计算,如何求它的导数,首先可以确定的是,一共有两个操作数,分别是 X 和 W 。
所以我们要分别求 Z 对 X 和 Z 对 W 的导数。我们先来求 Z 对 X 的导数。我们假设已经传到 Z 的梯度为 δ Z \delta_{Z} δ Z 。
则δ X \delta_{X} δ X 由下列项组成。
δ x 00 = W 00 δ z 00 δ x 01 = W 01 δ z 00 + W 00 δ z 01 δ x 02 = W 01 δ z 01 δ x 10 = W 10 δ z 00 + W 00 δ z 10 δ x 11 = W 11 δ z 00 + W 10 δ z 01 + W 01 δ z 10 + W 00 δ z 11 δ x 12 = W 11 δ z 01 + W 01 δ z 11 δ x 20 = W 10 δ z 10 δ x 21 = W 11 δ z 10 + W 10 δ z 11 δ x 22 = W 11 δ z 11 \begin{equation}
\begin{aligned}
\delta_{x00} &= W_{00}\delta_{z00} \\
\delta_{x01} &= W_{01}\delta_{z00} + W_{00}\delta_{z01} \\
\delta_{x02} &= W_{01}\delta_{z01} \\
\delta_{x10} &= W_{10}\delta_{z00} + W_{00}\delta_{z10} \\
\delta_{x11} &= W_{11}\delta_{z00} + W_{10}\delta_{z01} + W_{01}\delta_{z10} + W_{00}\delta_{z11} \\
\delta_{x12} &= W_{11}\delta_{z01} + W_{01}\delta_{z11} \\
\delta_{x20} &= W_{10}\delta_{z10} \\
\delta_{x21} &= W_{11}\delta_{z10} + W_{10}\delta_{z11} \\
\delta_{x22} &= W_{11}\delta_{z11} \\
\end{aligned}
\end{equation}
δ x 00 δ x 01 δ x 02 δ x 10 δ x 11 δ x 12 δ x 20 δ x 21 δ x 22 = W 00 δ z 00 = W 01 δ z 00 + W 00 δ z 01 = W 01 δ z 01 = W 10 δ z 00 + W 00 δ z 10 = W 11 δ z 00 + W 10 δ z 01 + W 01 δ z 10 + W 00 δ z 11 = W 11 δ z 01 + W 01 δ z 11 = W 10 δ z 10 = W 11 δ z 10 + W 10 δ z 11 = W 11 δ z 11
乍一看十分复杂,但事实上,上面的计算过程可以看成下面的一次卷积计算:
[ 0 0 0 0 0 δ z 00 δ z 01 0 0 δ z 10 δ z 11 0 0 0 0 0 ] ∗ [ W 11 W 10 W 01 W 00 ] = [ δ x 00 δ x 01 δ x 02 δ x 10 δ x 11 δ x 12 δ x 20 δ x 21 δ x 22 ] \begin{equation}
\begin{bmatrix}
0 & 0 & 0 & 0 \\
0 & \delta_{z00} & \delta_{z01} & 0 \\
0 & \delta_{z10} & \delta_{z11} & 0 \\
0 & 0 & 0 & 0 \\
\end{bmatrix}
*
\begin{bmatrix}
W_{11} & W_{10} \\
W_{01} & W_{00} \\
\end{bmatrix}
=
\begin{bmatrix}
\delta_{x00} & \delta_{x01} & \delta_{x02} \\
\delta_{x10} & \delta_{x11} & \delta_{x12} \\
\delta_{x20} & \delta_{x21} & \delta_{x22} \\
\end{bmatrix}
\end{equation}
0 0 0 0 0 δ z 00 δ z 10 0 0 δ z 01 δ z 11 0 0 0 0 0 ∗ [ W 11 W 01 W 10 W 00 ] = δ x 00 δ x 10 δ x 20 δ x 01 δ x 11 δ x 21 δ x 02 δ x 12 δ x 22
其中,左侧是 Z 的梯度矩阵 $ \delta_{Z} $ 经过一圈padding之后的样子,中间的卷积核是原来的卷积核 W ,经过180度的旋转所得到的。
这有些看起来不可思议,但确实就是这么巧妙,事实上这个操作又被称之为反卷积。再深入一点的数学原理,比较复杂,不好阐述,
我们就根据这个例子给出一个通用的公式:
δ X = pad ( dilate ( δ Z ) ) ∗ rotate180 ( W ) \begin{equation}
\delta_X = \text{pad}(\text{dilate}(\delta_Z)) \ast \text{rotate180}(W)
\end{equation}
δ X = pad ( dilate ( δ Z )) ∗ rotate180 ( W )
在等式(5)中,我们又多出来了一个 dilate 操作。这个操作是干什么的呢?当我们考虑正向卷积的 stride 大于1的情况,这时候会发现对于矩阵 Z 来说,输入矩阵 X 的贡献不是连续的。什么意思呢,你可以考虑一个输入矩阵为 4x4 的维度并且 stride 为2的情况。这时候 z01 的前向计算贡献显然与 z00 的贡献没有相同的元素。所以 delta_z00 和 delta_z01 肯定不是相邻的,而需要插入一些0。我们称这种操作为膨胀,即 dilate。我们可以这样理解,一个 z 是由卷积核走了一步得到的,那么它对应的那一步就要跨过 S(stride 的大小)个元素,所以排除掉 z 自身,就应该有 S - 1 个空位置在每个 z 之间。这样下一个 z 和前一个 z 的贡献来源才是正确的。dilate 等于 S - 1。
然后我们需要理解 padding 需要多大。这里有一个非常简单的方法,我们只需要确定 delta_z00 周围一圈填了几层0就行了。而 delta_z00 一定会由 W00 计算得到。一个在左上角,一个在右下角,结合卷积计算的原理,你只需要把卷积核放到这两个元素重合的位置,就能知道要 padding 的大小了,很显然,padding 等于 K - P - 1。这里的 P 指的是正向卷积时候的 padding,除去正向卷积一开始补的0,计算就和我们说的上述过程一样。
上述的理解过程,其实还可以通过三个矩阵的形状经过卷积计算要求匹配,来用公式推导 dilate 和 padding 的大小,我这里给出的是我自己的理解。
接下来我们来看如何求 δ W \delta_{W} δ W 。我们直接通过等式(2)给出每一项的计算公式。
δ w 00 = x 00 δ z 00 + x 01 δ z 01 + x 10 δ z 10 + x 11 δ z 11 δ w 01 = x 01 δ z 00 + x 02 δ z 01 + x 11 δ z 10 + x 12 δ z 11 δ w 10 = x 10 δ z 00 + x 11 δ z 01 + x 20 δ z 10 + x 21 δ z 11 δ w 11 = x 11 δ z 00 + x 12 δ z 01 + x 21 δ z 10 + x 22 δ z 11 \begin{equation}
\begin{aligned}
\delta_{w00} &= x_{00}\delta_{z00} + x_{01}\delta_{z01} + x_{10}\delta_{z10} + x_{11}\delta_{z11} \\
\delta_{w01} &= x_{01}\delta_{z00} + x_{02}\delta_{z01} + x_{11}\delta_{z10} + x_{12}\delta_{z11} \\
\delta_{w10} &= x_{10}\delta_{z00} + x_{11}\delta_{z01} + x_{20}\delta_{z10} + x_{21}\delta_{z11} \\
\delta_{w11} &= x_{11}\delta_{z00} + x_{12}\delta_{z01} + x_{21}\delta_{z10} + x_{22}\delta_{z11}
\end{aligned}
\end{equation}
δ w 00 δ w 01 δ w 10 δ w 11 = x 00 δ z 00 + x 01 δ z 01 + x 10 δ z 10 + x 11 δ z 11 = x 01 δ z 00 + x 02 δ z 01 + x 11 δ z 10 + x 12 δ z 11 = x 10 δ z 00 + x 11 δ z 01 + x 20 δ z 10 + x 21 δ z 11 = x 11 δ z 00 + x 12 δ z 01 + x 21 δ z 10 + x 22 δ z 11
不难发现,这是里面蕴含一个非常明显的卷积过程。我们直接给出卷积公式
[ X 00 X 01 X 02 X 10 X 11 X 12 X 20 X 21 X 22 ] ∗ [ δ z 00 δ z 01 δ z 10 δ z 11 ] = [ δ w 00 δ w 01 δ w 10 δ w 11 ] \begin{equation}
\begin{bmatrix}
X_{00} & X_{01} & X_{02} \\
X_{10} & X_{11} & X_{12} \\
X_{20} & X_{21} & X_{22}
\end{bmatrix}
*
\begin{bmatrix}
\delta_{z00} & \delta_{z01} \\
\delta_{z10} & \delta_{z11}
\end{bmatrix}
=
\begin{bmatrix}
\delta_{w00} & \delta_{w01} \\
\delta_{w10} & \delta_{w11}
\end{bmatrix}
\end{equation}
X 00 X 10 X 20 X 01 X 11 X 21 X 02 X 12 X 22 ∗ [ δ z 00 δ z 10 δ z 01 δ z 11 ] = [ δ w 00 δ w 10 δ w 01 δ w 11 ]
所以我们可以很轻松的给出最后的通用公式,值得注意的是,这里的 δ Z \delta_{Z} δ Z 也需要dilate操作,并且和我们上文讲的是同一计算方式。
δ W = X ∗ dilate ( δ Z ) \begin{equation}
\delta_W = X \ast \text{dilate}(\delta_Z)
\end{equation}
δ W = X ∗ dilate ( δ Z )
我们终于讲完了卷积反向传播的计算原理。我们现在来看如何通过代码实现高维的卷积反向传播,这其中又会涉及到输入通道和输出通道,这两个维度的变换。
我们直接给出代码,并在代码中详细解释。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 def gradient (self, out_grad: Value, node: Value ) -> Value | Tuple [Value]: input , weight = node.inputs grad_dilate = dilate(out_grad, (1 , 2 ), self.stride - 1 ) weight_r180 = flip(weight, (0 , 1 )) weight_t = transpose(weight_r180) K = weight_r180.shape[0 ] grad_input = conv(grad_dilate, weight_t, 1 , K - 1 - self.padding) grad_dilate = grad_dilate.transpose((0 , 2 )).transpose((0 , 1 )) input_t = transpose(input ,(0 ,3 )) grad_weight = conv(input_t, grad_dilate, 1 , self.padding) grad_weight = grad_weight.transpose((0 , 2 )).transpose((0 , 1 )) return Tensor(grad_input), grad_weight
实际上前面的这部分卷积的讲解,是我从我的某个大作业里copy过来的,反传计算公式是抄的一位老哥写的博客,但是我丰富了一些自己的理解进去,讲得更简单易懂。也是让自己稍微复习一下(现在又忘记了),接下来是分享我最近弄的一个非常简单的MaxPooling算子的img2col实现,
老规矩,我们直接给出实现的前向计算代码,我会在代码中详细讲解:
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 class MaxPooling2D (TensorOp ): def __init__ (self,kernel_size: int , stride: Optional [int ] = 1 , padding: Optional [int ] = 0 ): self.stride = stride self.padding = padding self.kernel_size = kernel_size self._A_precomputed = None def precompute (self, A ): N, H, W, C = A.shape K = self.kernel_size stride = self.stride padding = self.padding pad_width = [(0 , 0 ), (0 , 0 ), (padding, padding), (padding, padding)] _A = array_api.pad(A, pad_width, mode='constant' , constant_values=0 ) if self.padding > 0 else A H_out = (H - K + 2 * padding) // stride + 1 W_out = (W - K + 2 * padding) // stride + 1 Ns, Hs, Ws, Cs = _A.strides _A = array_api.as_strided( _A, shape=(N, H_out, W_out, K, K, C), strides=(Ns, Hs*stride, Ws*stride, Hs, Ws, Cs) ) self._A_precomputed = compact(_A) def compute (self, A ): self.precompute(A) max_pool_out = array_api.max (self._A_precomputed, axis=(3 , 4 ),keepdims=True ) return max_pool_out
看上去非常简单,非常的nice,ok我现在来想一下,如何对这个MaxPoolling算子求导,那么随之而来的第一个问题就是,max的导数怎么求?
其实这里有一个非常简单的理解:取了这个元素的位置就是1,没有取的位置就是0。这里我给一个我自己实现的max算子的计算和反向传播,可能会有助于理解,当然我们实际上没有用到这个算子。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 class Max (TensorOp ): def __init__ (self, axes: Optional [tuple ] = None ): self.axes = axes if isinstance (self.axes, int ): self.axes = tuple ([self.axes]) def compute (self, a ): return array_api.max (a, axis=self.axes, keepdims=True ) def gradient (self, out_grad, node ): a = node.inputs[0 ] array = a.cached_data max_values = self.compute(array) mask = (array== max_values) grad_sum = array_api.sum (mask, axis=self.axes, keepdims=True ) grad = out_grad * mask / grad_sum return grad
在这个基础上,对于一个池化算子,我们可以先把它在kenerl里面最大值的索引记录下来。
记录索引,这样我们才能知道对于输入的X来说哪些地方是要保留的,哪些地方可以去掉。 接着我们遍历所有的最大值索引位置,从输入X里面把对应的值取出来,累加到一个空矩阵里(这个矩阵和输入矩阵的维度一样,累加的原因是可能有重复的索引,贡献了多次)
这样我们就能得到这一层的max导数了。还有一些细节我们写到代码注释里。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 def gradient (self, out_grad: Tensor, node: Tensor ): array = self._A_precomputed array = array.transpose((0 ,1 ,2 ,5 ,4 ,3 )) N,H,W,C,K,_ = array.shape argmax = array_api.argmax(array.reshape(N,H,W,C,K**2 ), axes=-1 ) grad_input = array_api.zeros_like(node.inputs[0 ].cached_data) for i, max_index in array_api.ndenumerate(argmax): n,h,w,c = i mh, mw = max_index // K, max_index % K grad_input[n, h * self.stride + mh, w * self.stride + mw, c] += node.inputs[0 ].cached_data[i] grad_dilate = dilate(out_grad,(1 , 2 ), self.stride - 1 ) return summation(grad_dilate,(-3 ,-2 )) * grad_input
公式就是这样,大家也可以看公式理解,这里后面的偏导数就是max函数的求导
δ X = dilate ( δ Z ) ⊙ ∂ Z ∂ X \begin{equation}
\delta_X = \text{dilate}(\delta_Z) \odot \frac{\partial Z}{\partial X}
\end{equation}
δ X = dilate ( δ Z ) ⊙ ∂ X ∂ Z
本篇到此结束。之后可能更新一些transformer的原理实现。挖个坑先()。