pytorch中的torch.nn.Unfold和torch.nn.Fold

目的
- Unfold
- Fold

目的

平时使用卷积操作时，既卷积核滑动窗口操作，对于pytorch，以二维图像为例，调用nn.Conv2d就能完成对输入（feature maps）的卷积操作。
但有时，maybe要探究卷积核对应的某一channel的单个窗口的卷积操作，或显式地进行卷积操作。此时，就需要nn.Unfold和nn.Fold。前段时间引起较大争议的BagNet（Bag of local feature net） 的分块卷积操作既由此函数完成。
一般来说，Conv2d 就是 Unfold + matmul + fold

Unfold

torch.nn.Unfold按照官方的说法，既从一个batch的样本中，提取出滑动的局部区域块，也就是卷积操作中的提取kernel filter对应的滑动窗口。

如上图所示，蓝色框部分就是kernel filter（红色框部分）对应的滑动窗口。
首先来看下torch.nn.Unfold的参数：

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

跟nn.Conv2d的参数很相似，卷积核的尺寸，空洞大小，填充大小和步长。

官方解释中：unfold的输入为（NNN, CCC, HHH, WWW），其中N为batch_size，C是channel个数，H和W分别是channel的长宽。则unfold的输出为（NNN, C×∏C \times \prodC×∏(kernel_size), LLL），其中∏\prod∏(kernel_size)为kernel_size长和宽的乘积， L是channel的长宽根据kernel_size的长宽滑动裁剪后，得到的区块的数量。

以输入（1, 3, 4, 4）为例，设定kernel_size = (2, 2)，stride = 2，根据官方给出的LLL计算公式：L=∏d⌊spatial_size[d]+2×padding[d]−dilation[d]×(kernel_size[d]−1)stride[d]+1⌋L = \prod_d \lfloor{\dfrac {spatial\_size[d] + 2 \times padding[d] - dilation[d] \times (kernel\_size[d]-1)}{stride[d]} + 1}\rfloorL=d∏⌊stride[d]spatial_size[d]+2×padding[d]−dilation[d]×(kernel_size[d]−1)+1⌋ddd是channel的维度，二维图像既长宽的维度。
则LLL（区块数量）为⌊(4+2×0−1×(2−1)−12+1)⌋×⌊(4+2×0−1×(2−1)−12+1)⌋=4\lfloor(\dfrac{4 + 2 \times 0 - 1 \times (2-1) -1}{2} + 1) \rfloor \times \lfloor(\dfrac{4 + 2 \times 0 - 1 \times (2-1) -1}{2} + 1) \rfloor = 4⌊(24+2×0−1×(2−1)−1+1)⌋×⌊(24+2×0−1×(2−1)−1+1)⌋=4，每个区块的大小为C×kernel_size[0]×kernel_size[1]C \times kernel\_size[0] \times kernel\_size[1]C×kernel_size[0]×kernel_size[1]，既2×2×2=82 \times 2 \times 2 = 82×2×2=8，做为输出的第二个维度。
为了更直观的展示unfold函数所做的操作，以下述代码和结果为例：

inputs = torch.randn(1, 2, 4, 4)
print(inputs.size())
print(inputs)
unfold = torch.nn.Unfold(kernel_size=(2, 2), stride=2)
patches = unfold(inputs)
print(patches.size())
print(patches)

torch.Size([1, 2, 4, 4])
tensor([[[[ 1.4818, -0.1026, -1.7688,  0.5384],[-0.4693, -0.0775, -0.7504,  0.2283],[-0.1414,  1.0006, -0.0942,  2.2981],[-0.9429,  1.1908,  0.9374, -1.3168]],[[-1.8184, -0.3926,  0.1875,  1.3847],[-0.4124,  0.9766, -1.3303, -0.0970],[ 1.7679,  0.6961, -1.6445,  0.7482],[ 0.1729, -0.3196, -0.1528,  0.2180]]]])
torch.Size([1, 8, 4])
tensor([[[ 1.4818, -1.7688, -0.1414, -0.0942],[-0.1026,  0.5384,  1.0006,  2.2981],[-0.4693, -0.7504, -0.9429,  0.9374],[-0.0775,  0.2283,  1.1908, -1.3168],[-1.8184,  0.1875,  1.7679, -1.6445],[-0.3926,  1.3847,  0.6961,  0.7482],[-0.4124, -1.3303,  0.1729, -0.1528],[ 0.9766, -0.0970, -0.3196,  0.2180]]])

对代码结果分析，nn.Unfold对输入channel的每一个kernel_size[0]×kernel_size[1]kernel\_size[0] \times kernel\_size[1]kernel_size[0]×kernel_size[1]的滑动窗口区块做了展平操作。

Fold

torch.nn.Fold的操作与Unfold相反，将提取出的滑动局部区域块还原成batch的张量形式。
代码如下：

fold = torch.nn.Fold(output_size=(4, 4), kernel_size=(2, 2), stride=2)
inputs_restore = fold(patches)
print(inputs_restore)
print(inputs_restore.size())

tensor([[[[ 1.4818, -0.1026, -1.7688,  0.5384],[-0.4693, -0.0775, -0.7504,  0.2283],[-0.1414,  1.0006, -0.0942,  2.2981],[-0.9429,  1.1908,  0.9374, -1.3168]],[[-1.8184, -0.3926,  0.1875,  1.3847],[-0.4124,  0.9766, -1.3303, -0.0970],[ 1.7679,  0.6961, -1.6445,  0.7482],[ 0.1729, -0.3196, -0.1528,  0.2180]]]])
torch.Size([1, 2, 4, 4])

分析结果，Fold的操作通过设定output_size=(4, 4)，完成与Unfold的互逆的操作。