自己实现一个Pytorch二维卷积

一.前言

卷积是卷积神经网络的核心，其主要包括一维卷积、二维卷积和三维卷积，其中二维卷积是应用最广泛的，例如我们常见的图像上的卷积。今天博主就来带大家手撕一下二维卷积计算的实现，注意本文专注于卷积运算的实现，话不多说，请看下文。

二.预备知识

注：限于篇幅关系，暂不考虑stride和padding。

先给个栗子吧，假设存在一张大小为$3\times 4\times 4$​图像，即图像存在三个通道 (channel)，每个channel的宽高都为$4$​​​​​，即图1中的3个深蓝色模块。​我们知道若要在该图像上做卷积，则必须要定义一个3通道的卷积核，即图中的3个绿色模块。

在进行卷积运算的过程中，图像的每个通道都会和与之对应的卷积核进行卷积运算，具体为卷积核在单个通道上按从上到下、从左往右的顺序滑动，每到一个位置便与卷积核对应位置的权重相乘再求和（例如通道1中橙色位置运算与卷积核对应位置的元素相乘再求和得到输出结果中紫色方块中的值）。三个通道都分别与对应的卷积核进行运算，得到3个输出，然后3个输出沿通道方向进行求和，得到最终的结果，即图中的淡橙色模块。

图1: 二维卷积运算示例

当然，该栗子中只演示了单个卷积核与图像进行卷积运算，事实上实际应用中通常通过不同的卷积核来提取不同的特征，即不同的卷积核提取的特征（淡橙色的模块）是不同的。

三.暴力实现卷积

按照第二节的思路，我们很容易就可以暴力实现二维卷积运算，示例代码如下所示：

def conv_on_schannel(X,K):
    """
    功能：单通道上的卷积
    x: 输入， shape (H,W)
    K: 卷积核，shape (k_h,k_w)
    """
    h,w = K.shape
    Y = torch.zeros(X.shape[0] - h + 1, X.shape[1] - w + 1)
    # 在单通道上滑动
    for i in range(Y.shape[0]):
        for j in range(Y.shape[1]):
            Y[i,j] = (X[i:i + h,j:j + w] * K).sum()
    
    return Y

def conv_on_mchannel(X,K):
    """
    功能：单个卷积核在feature map上做卷积运算
    X: feature map，shape (C,H,W)
    K: kernel，shape (C,k_h,k_w)
    """
    res = conv_on_schannel(X[0,:,:],K[0,:,:])
    for i in range(1, X.shape[0]):
        res += conv_on_schannel(X[i, :, :], K[i, :, :])
    return res

def conv2D(X, K):
    """
    功能：多个卷积核在feature map上做卷积运算
    x: feature map, shape (C,H,W), C表示通道数
    K: multi kernel, shape (k_n,C,k_h,k_w), k_n表示卷积核的个数
    """
    return torch.stack([conv_on_mchannel(X, k) for k in K])

# 图1中的图像
x = torch.tensor([
    [[2,1,0,2],[4,2,1,1],[0,3,2,2],[4,1,5,1]],
    [[1,0,0,2],[4,5,3,0],[2,2,2,2],[1,3,3,1]],
    [[1,2,2,1],[0,4,5,1],[4,6,8,1],[0,1,5,1]]
], dtype=torch.float32)
# 图1中的卷积核
k = torch.tensor([
    [[1,-1],[1,-1]],
    [[1,1],[-1,-1]],
    [[1,-1],[-1,1]]
], dtype=torch.float32)
# 扩充卷积核数量维度
k = k.unsqueeze(0)
print(conv2D(x,k))
"""
tensor([[[-2., -5., -6.],
         [ 2.,  7., -4.],
         [-1., -3.,  7.]]])
"""

该种实现方式可读性很强，但是时间复杂度太高，对于一张图像需要先遍历一遍卷积核，然后遍历一遍图像的每个通道，然后还需要在通道上进行窗口滑动（又是二重循环）。做个测试，预定义一张大小为$3\times 200\times 200$​​​​的图像，然后用两个kernel size为$2\times 2$的卷积核与其进行卷积运算，即：

# 通道为3，宽高都为200
x = torch.rand(3,200,200)
# 2个卷积核
k = torch.rand(2,3,2,2)
starttime = datetime.datetime.now()
Y = conv2D(x,k)
endtime = datetime.datetime.now()
print((endtime - starttime).seconds)

在博主的PC上测试，在其上完成卷积运算需要大概3秒，而Pytorch自带的torch.nn.functional.conv2d函数进行同样的操作耗时却几乎为0秒，显然这不是我们所需要的是实现方式。

四.优雅实现卷积

根据第三节的介绍，我们需要一种时间复杂度更低的实现方式，很容易我们可以想到能不能让多个卷积核同时在图像的多个通道上进行卷积运算，这就需要矩阵乘法来大显身手了。再回头看看图1，第一节介绍过，卷积核的三个通道在图像的三个通道上是同时进行滑动的，因此我们只需要将当前时刻图像滑动的位置的元素取出来进行展平，同时将卷积核的三个通道的元素进行凭借展平，然后计算点积，即可获得最终结果中的一个元素，示例参见图2：

图2: 单次多通道同时卷积

图2中左侧为从图像的3个通道取出的卷积核出现位置的对应元素拼接成的向量，右侧有展平的卷积核，同颜色表示一一对应，然后左侧向量与右侧向量进行点积，即:
$$2\times 1+1\times -1+4\times 1+2\times -1+1\times 1+0\times 1+4\times -1+5\times -1+1\times 1+2\times -1+0\times -1+4\times 1=-2$$
卷积核进行展平很容易实现，关键是如何快速寻找卷积核在图像上移动时，图像上卷积核位置的元素。其实torch针对该功能已经实现好了优化过的API，即:

torch.nn.functional.unfold(input, kernel_size)

该通过unfold函数，我们可以获取图像上对应卷积核滑动位置的元素凭借成的类似图2中的向量，示例代码如下：

# input必须为4D Tensor
print(F.unfold(x.unsqueeze(0),kernel_size=(2,2)))
"""
tensor([[[2., 1., 0., 4., 2., 1., 0., 3., 2.],
         [1., 0., 2., 2., 1., 1., 3., 2., 2.],
         [4., 2., 1., 0., 3., 2., 4., 1., 5.],
         [2., 1., 1., 3., 2., 2., 1., 5., 1.],
         [1., 0., 0., 4., 5., 3., 2., 2., 2.],
         [0., 0., 2., 5., 3., 0., 2., 2., 2.],
         [4., 5., 3., 2., 2., 2., 1., 3., 3.],
         [5., 3., 0., 2., 2., 2., 3., 3., 1.],
         [1., 2., 2., 0., 4., 5., 4., 6., 8.],
         [2., 2., 1., 4., 5., 1., 6., 8., 1.],
         [0., 4., 5., 4., 6., 8., 0., 1., 5.],
         [4., 5., 1., 6., 8., 1., 1., 5., 1.]]])
"""

可以看出输出矩阵中第3维的每一列就是我们所要寻找的类似图2左侧的向量。然后，我们对卷积核进行展平 (图2右侧向量)，即：

print(k.flatten())
"""
tensor([ 1, -1,  1, -1,  1,  1, -1, -1,  1, -1, -1,  1])
"""

将二者进行矩阵相乘即可完成多个卷积核在图像上的卷积运算，下面给出实现代码：

def conv2D_V1(X, K):
    """
    功能：多个卷积核在feature map上卷积
    x: feature map, shape (N,C,H,W), N表示图片数量,C表示图片的通道数
    K: multi kernel, shape (k_n,C,k_h,k_w)
    """
    H,W = X.shape[2],X.shape[3]
    k_h,k_w = K.shape[2],K.shape[3]
    # 获取X上的所有滑动窗口
    X_uf = F.unfold(X,kernel_size=K.shape[2:])
    # 卷积
    corr_uf = X_uf.transpose(1,2).matmul(K.view(K.shape[0],-1).T).transpose(1,2)
    # 计算卷积后的feature map的size
    h_n,w_n = H - k_h + 1,W - k_w + 1
    Y = F.fold(corr_uf,output_size=(h_n,w_n),kernel_size=(1,1))
    
    return Y

print(conv2D_V1(x.unsqueeze(0), k.unsqueeze(0)))
"""
tensor([[[[-2., -5., -6.],
          [ 2.,  7., -4.],
          [-1., -3.,  7.]]]])
"""

在该版本的函数中还实现了多张图片的同时卷积，因此输入的X也要为四位张量，即(图片数, 图片通道数, 高、宽)。采用同暴力实现相同的测试用例，发现使用conv2D_V1进行卷积的时间几乎为0秒。

五.结语

以上便是本文的全部内容，要是觉得不错的话就点个赞或关注一下博主吧，你们的支持是博主继续创作的不解动力，当然若是有任何问题也敬请批评指正！！！