引言:一次小小的复习

卷积和池化的 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):
# 初始化stride和padding
def __init__(self, stride: Optional[int] = 1, padding: Optional[int] = 0):
self.stride = stride
self.padding = padding

def compute(self, A, B):
# 取出输入A的四个维度,N代表批量,和pytorch定义的输入维度:N C H W类似,但是顺序不同。
N, H, W, C_in = A.shape
# 取出卷积核的四个维度(这里默认高和宽都是K),I是输入通道数,这里为了判断操作是否合法,起了别名
K, _, I, C_out = B.shape
assert C_in == I, "input tensor shape and kernel doesn't match"
# 只用对高和宽这两个维度padding
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
# 得到每个轴的stride,这么做的原因是为了用as_strided方法来快速的构建img2col中的大矩阵。
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)
)
# compact操作是复制一份对象变成紧凑数组,因为as_strided这个方法只是取出了地址索引,实际存储是离散的,这样计算速度会下降。
# 在构建完大矩阵后,就是col的部分,变成向量点积
_A_ = compact(_A).reshape((-1, inner_dim))
_B_ = compact(B).reshape((-1, C_out))
out = _A_ @ _B_
# 最后把输出重塑成我们需要的维度:N H W C
return out.reshape((N, H_out, W_out, C_out))

现在我们重点关注卷积反向传播部分的代码,如何高效的实现卷积的反向传播。为了便于理解,我们从一个简单的例子开始。考虑如下的卷积操作:

[X00X01X02X10X11X12X20X21X22][W00W01W10W11]=[Z00Z01Z10Z11]\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 矩阵作为我们的输入,W 矩阵作为我们的卷积核(参数矩阵),Z 矩阵作为我们的输出矩阵。没有padding,
stride默认为1,输入输出通道数都是1。于是我们可以得到如下的计算公式。

{Z00=X00W00+X01W01+X10W10+X11W11Z01=X01W00+X02W01+X11W10+X12W11Z10=X10W00+X11W01+X20W10+X21W11Z11=X11W00+X12W01+X21W10+X22W11\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}

现在我们思考对于这样的一个卷积计算,如何求它的导数,首先可以确定的是,一共有两个操作数,分别是 X 和 W 。
所以我们要分别求 Z 对 X 和 Z 对 W 的导数。我们先来求 Z 对 X 的导数。我们假设已经传到 Z 的梯度为 δZ\delta_{Z}
δX\delta_{X} 由下列项组成。

δx00=W00δz00δx01=W01δz00+W00δz01δx02=W01δz01δx10=W10δz00+W00δz10δx11=W11δz00+W10δz01+W01δz10+W00δz11δx12=W11δz01+W01δz11δx20=W10δz10δx21=W11δz10+W10δz11δx22=W11δz11\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}

乍一看十分复杂,但事实上,上面的计算过程可以看成下面的一次卷积计算:

[00000δz00δz0100δz10δz1100000][W11W10W01W00]=[δx00δx01δx02δx10δx11δx12δx20δx21δx22]\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}

其中,左侧是 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}

在等式(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}。我们直接通过等式(2)给出每一项的计算公式。

δw00=x00δz00+x01δz01+x10δz10+x11δz11δw01=x01δz00+x02δz01+x11δz10+x12δz11δw10=x10δz00+x11δz01+x20δz10+x21δz11δw11=x11δz00+x12δz01+x21δz10+x22δz11\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}

不难发现,这是里面蕴含一个非常明显的卷积过程。我们直接给出卷积公式

[X00X01X02X10X11X12X20X21X22][δz00δz01δz10δz11]=[δw00δw01δw10δw11]\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}

所以我们可以很轻松的给出最后的通用公式,值得注意的是,这里的 δZ\delta_{Z} 也需要dilate操作,并且和我们上文讲的是同一计算方式。

δW=Xdilate(δZ)\begin{equation} \delta_W = X \ast \text{dilate}(\delta_Z) \end{equation}

我们终于讲完了卷积反向传播的计算原理。我们现在来看如何通过代码实现高维的卷积反向传播,这其中又会涉及到输入通道和输出通道,这两个维度的变换。
我们直接给出代码,并在代码中详细解释。

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
# out_grad 就是 delta_z , node就是当前节点,它负责记录输入:X,W和节点的计算操作:conv
# 这是Conv类里的一个函数
def gradient(self, out_grad: Value, node: Value) -> Value | Tuple[Value]:
input, weight = node.inputs

# 对input的导数
# 根据原理图,需要padding一圈,但是我们的conv自带了这个操作,所以直接考虑dilate操作
# N H W C 所以是 1, 2 轴需要dilate, dilate = s-1
grad_dilate = dilate(out_grad, (1, 2), self.stride - 1)
# 然后翻转180度,就是沿着K_h和K_w方向分别翻转一次。K K C_in C_out
weight_r180 = flip(weight, (0, 1))
# 注意为了保证梯度的形状和输入一样,所以最后卷积输出的通道数应该是C_in,所以我们调换一下通道顺序
weight_t = transpose(weight_r180)
K = weight_r180.shape[0]
# 对X的梯度,我们在卷积的时候进行padding操作, padding = K - P - 1
grad_input = conv(grad_dilate, weight_t, 1, K - 1 - self.padding)

# 下面求对 W 的导数
# 这里非常巧妙,我们可以想一想如何得到W的形状: K K C_in C_out
# 我们希望保留 C_in和 C_out这两个参数,把 N 给去掉
# 我们可以把 N 直接放到 C_in 原本的位置上,这样卷积的时候就会去掉这个参数
# 然后我们还需要对 grad_dilate再换一次轴,让它看起来像一个卷积核 ,H W两个维度在前面
grad_dilate = grad_dilate.transpose((0, 2)).transpose((0, 1))
# 这里也是为了消 N
input_t = transpose(input,(0,3))
# 直接卷积计算
grad_weight = conv(input_t, grad_dilate, 1, self.padding)
# 最后得到是一个 C_in H W C_out 的东西,所以要把轴再换回来
grad_weight = grad_weight.transpose((0, 2)).transpose((0, 1))
# grad_input不是tensor,所以需要变成tensor(经过了tensor计算的数组会自动变成tensor,这是我们在构造计算图时实现的)
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):
# 初始化stride和padding,kernel_size
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
# 多了一个字段,用来存储img2col后的矩阵
self._A_precomputed = None

# 这里的我的实现其实比较ugly,没有很好的抽象和复用,又写了一遍img2col
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)
# 计算每个窗口的最大值
# 注意这里的输入形状是 N, H_out, W_out, K, K, C,所以要求每个窗口的最大值就是 kernel所处的位置 3,4轴
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] # 假设node是一个封装了数据的对象
array = a.cached_data
# 计算最大值,保持维度以便广播操作
max_values = self.compute(array)
# 生成一个布尔数组,其中最大值位置为True
mask = (array== max_values)
# 计算输入梯度的总和
grad_sum = array_api.sum(mask, axis=self.axes, keepdims=True)
# 将外部梯度(out_grad)广播到与输入数据相同的形状,然后乘以mask。
# 这样,只有最大值位置才会获得梯度值。最后,我们需要将梯度值平均分配到所有最大值位置,
# 这是通过除以grad_sum实现的,以处理存在多个最大值的情况。
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):
# 获取img2col矩阵
array = self._A_precomputed
# 这里想要求我们所说的最大值索引,但是呢我们现在的维度有些不方便
# 所以把 N,H,W,K,K,C => N,H,W,C,K,K
# 这样的话我们可以很轻松的对最后两个维度进行合并,然后取最大值索引。
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)
# 生成一个和输入矩阵形状一样的全0矩阵
grad_input = array_api.zeros_like(node.inputs[0].cached_data)
# i的含义 max_index的含义
# 取出最大值所在的索引值 ,在维度 N H W C 中的索引 , 取出 k * k 这一个维度中的索引
for i, max_index in array_api.ndenumerate(argmax):

n,h,w,c = i
# 得到窗口 (k ,k)中实际的索引,行和列索引
mh, mw = max_index // K, max_index % K
# 对于一个输出的索引值来说,想要找到他的输入对应的位置,我们可以模拟这个窗口滑动的过程,
# 窗口与输入矩阵重叠的时候就可以利用最大值在k中的相对位置推导出在输入中的绝对位置,
# 想不明白的同学可以自己画一下
grad_input[n, h * self.stride + mh, w * self.stride + mw, c] += node.inputs[0].cached_data[i]
# 最后一步也是最容易忽视的一步,我们的outgrad的形状是池化后的形状,和当前的输入形状不匹配,
# 所以我们可以效仿卷积里dilate的做法,把outgrad膨胀的和当前的输入矩阵形状一致。
grad_dilate = dilate(out_grad,(1, 2), self.stride - 1)
# 最后我们对这个矩阵在 kerel_size的维度进行压缩,然后族元素乘就行
return summation(grad_dilate,(-3,-2)) * grad_input

# 补充一下node.inputs是返回输入元素,这里的池化算子是一个一元操作符,只有一个元素,所以直接用【0】访问
# cached_data是我封装的Tensor类的底层数组,类似numpy

公式就是这样,大家也可以看公式理解,这里后面的偏导数就是max函数的求导

δX=dilate(δZ)ZX\begin{equation} \delta_X = \text{dilate}(\delta_Z) \odot \frac{\partial Z}{\partial X} \end{equation}

本篇到此结束。之后可能更新一些transformer的原理实现。挖个坑先()。