详解pytorch fold和unfold用法

先上结论，conv = unfold + matmul + fold. 即卷积操作等价于，先unfold(展开)，再执行矩阵乘法matmul，然后再fold(折叠)。具体过程如下：

unfold函数将一个输入Tensor(N,C,H,W) 展开成 (N,C * K1 * K2, Blocks)，其中kernel形状为(K1,K2)，总的Block数为Blocks。即把输入的Tensor根据kernel的大小展开成Blocks个向量。Block的计算公式如下：
$$Blocks={\text{H}}_{blocks}\times W_{blocks}$$
其中:
$$H_{blocks}=\frac{H+2\ast padding[0]-kernel[0]}{stride[0]}+1$$

$$W_{blocks}=\frac{W+2\ast padding[1]-kernel[1]}{stride[1]}+1$$

代码举例：

inp = torch.randn(1, 3, 10, 12)
w = torch.randn(2, 3, 4, 5)
inp_unf = torch.nn.functional.unfold(inp, (4, 5))#shape of inp_unf is (1,3*4*5,7*8)

其中，inp_unf的shape计算过程如下
$$H_{blocks}=\frac{10-4}{1}+1=7$$

$$W_{blocks}=\frac{12-5}{1}+1=8$$

out_unf = inp_unf.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)
#shape of out_unf is (1,2,56)

以上代码相当于 inp_unf(1, 60, 56) .t() * w(2 , 3 * 4 * 5).t() → out_unf (1, 56, 2 ) → out_unf (1, 2, 56)

unfold + matmul已经完成，最后是 fold过程. fold过程其实就是unfold的反过程，即把向量折叠回矩阵形式。

out = torch.nn.functional.fold(out_unf, (7, 8), (1, 1))
#out.size() = (1,2,7,8)

以上过程其实等价于直接进行Conv，因此

(torch.nn.functional.conv2d(inp, w) - out).abs().max()
#tensor(1.9073e-06)

可以看出卷积的结果和经过了unfold + matmul + fold的结果差距为10的-6次方，几乎可以认为是相等的了。

总结

利用pytorch 中fold 和unfold的组合可以实现类似Conv操作的滑动窗口，其中如果同一个图片的每个block的参数都是相同的，那么称为参数共享，就是标准的卷积层；如果每个block的参数都不一样，那么就不是参数共享的，此时一般称为局部连接层(Local connected layer)。

参考

https://pytorch.org/docs/stable/generated/torch.nn.Unfold.html

https://pytorch.org/docs/stable/generated/torch.nn.Fold.html

https://blog.youkuaiyun.com/LoseInVain/article/details/88139435