nn.Unfold 函数 与 卷积计算的实现过程

官方文档

将4-D 输入张量 (batched image-like tensors)沿着通道维度取出数据

def unfold(input, kernel_size, dilation=1, padding=0, stride=1):  

import torch
import torch.nn as nn
unfold = nn.Unfold(kernel_size=(2, 3))
input = torch.randn(2, 5, 3, 4)
output = unfold(input)
# each patch contains 30 values (2x3=6 vectors, each of 5 channels)
# 4 blocks (2x3 kernels) in total in the 3x4 input
print(output.size()) # torch.Size([2, 30, 4])

可视化一下代码

import numpy
import torch
import torch.nn as nn
unfold = nn.Unfold(kernel_size=(2, 3))
input = torch.randn(2, 5, 4, 4)
input = torch.range(1,2*5*4*4).reshape(2,5,4,4)
output = unfold(input)

卷积实现的解释：Convolution = Unfold + Matrix Multiplication + Fold（即卷积是通过Unfold 后进行矩阵乘法再fold回去）

# Convolution is equivalent with Unfold + Matrix Multiplication + Fold (or view to output shape)
inp = torch.randn(1, 3, 10, 12)
w = torch.randn(2, 3, 4, 5)
inp_unf = torch.nn.functional.unfold(inp, (4, 5))
out_unf = inp_unf.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)
out = torch.nn.functional.fold(out_unf, (7, 8), (1, 1))
# or equivalently (and avoiding a copy),
# out = out_unf.view(1, 2, 7, 8)
(torch.nn.functional.conv2d(inp, w) - out).abs().max()