Pytorch 基于im2col手动实现卷积conv2d(基于nn.Unfold实现卷积)(向量内积实现)

      如果老老实实地实现卷积运算，估计要重复好几层的for语句。这样的实现有点麻烦，而且， NumPy中存在使用for语句后处理变慢的缺点（NumPy中，访问元素时最好不要用 for语句）

       如上图所示，我们每次取的input，我们可以把它拉直，拉成一个行向量。它跟kernel进行元素相乘再求和，就相当于这个输入行向量，再跟kernel的列向量(将kernel拉成列向量)进行相乘

      在上图中，为了便于观察，将步幅设置得很大，以使滤波器的应用区域不重叠。而在实际的卷积运算中，滤波器的应用区域几乎都是重叠的。在滤波器的应用区域重叠的情况下，使用im2col展开后，展开后的元素个数会多于原方块的元素个数。因此，使用im2col的实现存在比普通的实现消耗更多内存的缺点。但是，汇总成一个大的矩阵进行计算，对计算机的计算颇有益处。比如，在矩阵计算的库（线性代数库）等中，矩阵计算的实现已被高度最优化，可以高速地进行大矩阵的乘法运算。因此，通过归结到矩阵计算上，可以有效地利用线性代数库。

使用 im2col展开输入数据后，之后就只需将卷积层的kernel纵向展开为1列，并计算2个矩阵的乘积即可，如下图。

将矩阵转为列向量

x.reshape(-1,1)

代码

import torch
from torch import nn
import torch.nn.functional as F
import math

def im2col(img, kernel_h, kernel_w, stride=1):
    N, C, H, W = img.shape
    out_h = (H - kernel_h)//stride + 1
    out_w = (W - kernel_w)//stride + 1
 
    col = torch.zeros((N, C, kernel_h, kernel_w, out_h, out_w))
 
    for y in range(kernel_h):
        y_max = y + stride*out_h
        for x in range(kernel_w):
            x_max = x + stride*out_w
            col[:, :, y, x, :, :] = img[:, :, y:y_max:stride, x:x_max:stride]
 
    col = col.permute(0, 4, 5, 1, 2, 3).contiguous().reshape(N*out_h*out_w, -1)
    return col

def my_conv(input, kernel, stride=1, padding=0, bias=0):
    if padding > 0:
        input = F.pad(input, (padding,padding,padding,padding))
    batch_size = input.shape[0]
    input_h, input_w = input.shape[2:4]
    kernel_h, kernel_w = kernel.shape[2:4]
    out_channel, in_channel = kernel.shape[0:2]
    output_h = math.floor((input_h - kernel_h) / stride + 1)
    output_w = math.floor((input_w - kernel_w) / stride + 1)
    
    input_vector = im2col(input, kernel_h, kernel_w, stride=stride)
    kernel_vector = kernel.reshape(kernel.shape[0], -1).T
    output = input_vector @ kernel_vector + bias
    output = output.reshape(batch_size, output_h, output_w, out_channel).permute(0,3,1,2).contiguous()    
    #注意可不能写成下面这样
    # output = output.reshape(batch_size, out_channel, output_h, output_w)
    
    
    return output
 
batch_size = 4
in_channel = 3
out_channel = 16
input = torch.rand(batch_size, in_channel ,5,5)
kernel = torch.rand(out_channel, in_channel, 3,3)
bias = torch.rand(out_channel)
 
my_output = my_conv(input, kernel, padding=1, stride=2, bias=bias)
 
output = F.conv2d(input, kernel, padding=1, stride=2, bias=bias)
 
assert torch.allclose(my_output, output)

用nn.Unfold实现

import torch
from torch import nn
import torch.nn.functional as F
import math

def my_conv(input, kernel, stride=1, padding=0, bias=0):
    if padding > 0:
        input = F.pad(input, (padding,padding,padding,padding))
    batch_size = input.shape[0]
    input_h, input_w = input.shape[2:4]
    kernel_h, kernel_w = kernel.shape[2:4]
    out_channel, in_channel = kernel.shape[0:2]
    output_h = math.floor((input_h - kernel_h) / stride + 1)
    output_w = math.floor((input_w - kernel_w) / stride + 1)
    
    unfold = nn.Unfold(kernel_size=(kernel_h, kernel_w), stride=stride)
    input_vector = unfold(input)
    
    kernel_vector = kernel.reshape(kernel.shape[0], -1).T
    output = (input_vector.permute(0,2,1).contiguous() @ kernel_vector ) + bias
    output = output.reshape(batch_size, output_h, output_w, out_channel).permute(0,3,1,2).contiguous()    
    
    #注意可不能写成下面这样
    # output = output.reshape(batch_size, out_channel, output_h, output_w)
    
    
    return output
 
batch_size = 4
in_channel = 3
out_channel = 16
input = torch.rand(batch_size, in_channel ,5,5)
kernel = torch.rand(out_channel, in_channel, 3,3)
bias = torch.rand(out_channel)
 
my_output = my_conv(input, kernel, padding=1, stride=2, bias=bias)
 
output = F.conv2d(input, kernel, padding=1, stride=2, bias=bias)
 
assert torch.allclose(my_output, output)