pytorch中的torch.nn.Unfold和torch.nn.Fold
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的输入为（ N N N, C C C, H H H, W W W），其中N为batch_size，C是channel个数，H和W分别是channel的长宽。则unfold的输出为（ N N N, C × ∏ C \times \prod C×∏(kernel_size), L L L），其中 ∏ \prod ∏(kernel_size)为kernel_size长和宽的乘积， L是channel的长宽根据kernel_size的长宽滑动裁剪后，得到的区块的数量。
以输入（1, 3, 4, 4）为例，设定kernel_size = (2, 2)，stride = 2，根据官方给出的 L L L计算公式： L = ∏ d ⌊ s p a t i a l _ s i z e [ d ] + 2 × p a d d i n g [ d ] − d i l a t i o n [ d ] × ( k e r n e l _ s i z e [ d ] − 1 ) s t r i d e [ d ] + 1 ⌋ L = \prod_d \lfloor{\dfrac {spatial\_size[d] + 2 \times padding[d] - dilation[d] \times (kernel\_size[d]-1)}{stride[d]} + 1}\rfloor L=d∏⌊stride[d]spatial_size[d]+2×padding[d]−dilation[d]×(kernel_size[d]−1)+1⌋ d d d是channel的维度，二维图像既长宽的维度。
则 L L L（区块数量）为 ⌊ ( 4 + 2 × 0 − 1 × ( 2 − 1 ) − 1 2 + 1 ) ⌋ × ⌊ ( 4 + 2 × 0 − 1 × ( 2 − 1 ) − 1 2 + 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 × k e r n e l _ s i z e [ 0 ] × k e r n e l _ s i z e [ 1 ] C \times kernel\_size[0] \times kernel\_size[1] C×kernel_size[0]×kernel_size[1]，既 2 × 2 × 2 = 8 2 \times 2 \times 2 = 8 2×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的每一个 k e r n e l _ s i z e [ 0 ] × k e r n e l _ s i z e [ 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的互逆的操作。