在pytorch中,我们常用的卷积是封装好了的卷积,如nn.Conv2d, 对应到原理图的左上那个子图

但如果封装好的这个卷积操作不能满足我们想要更细粒度的操作的话,pytorch还为我们提供了 unfold , matmul , fold 三个操作（conv = unfold + matmul + fold）

torch.nn.Unfold就是原理图下面中间的那个,也就算把一个立体的tensor(feature)分成w_1*h_1个部分(kernel_size-sized block),然后把每一个准备和kernel相乘的部分拉直。该类的构造器的参数有：

torch.nn.Unfold(kernel_size, dilation=1, padding=0, stride=1)

我们来看下unfold的输入和输出，其输入形状如:(c_0,w_0,h_0), 输出就是(c_0*w_k*h_k,w_1*h_1)

unfold之后,我们构造c_1个可以学习的tensor(c_0,w_k,h_k)作为kernels,并把它像左下图那样展开成(c_1, c_0*w_k * h_k) ,注意这里的c_1是kernel的个数

然后就用pytorch自带的matmul,把kernels 的展开乘unfold 之后的input tensor(c_0*w_k*h_k,w_1*h_1)得到Output Maps,维度为(c_1,w_1*h_1)
————————————————

原文链接：https://blog.csdn.net/u010087338/article/details/113666140

用nn.Unfold实现

Pytorch 基于im2col手动实现卷积conv2d(基于nn.Unfold实现卷积)(向量内积实现)_hxxjxw的博客-CSDN博客_conv2d实现

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)