torch.nn.Unfold函数进行图像的patch块提取

在项目中要用到图像转为一个个的n*n的patch块，本以为可以白嫖transformer，但妹想到它是直接nn.Conv2d(in_chans, in_chans, kernel_size=patch_size, stride=patch_size)，合着是直接每个patch都投影成一个数值了
但苦手了一天，终于还是被我用unfold搞好了。

首先简单的定义x的维度为(b, c, h, w)->(1, 1, 4, 4)

x = torch.randn(1, 1, 4, 4)
x->tensor([[[
[-1.1258, -1.1524, -0.2506, -0.4339],
[ 0.8487,  0.6920, -0.3160, -2.1152],
[ 0.3223, -1.2633,  0.3500,  0.3081],
[ 0.1198,  1.2377,  1.1168, -0.2473]]]])

patch_size=n #你想设定的每个patch的大小
y = nn.Unfold(kernel_size=patch_size, stride=patch_size)(x)
# y.shape->(b, c*k*k, patch_number) 这是定义
#可以看出x中每个patch在y中都变成了一列
tensor([[
[-1.1258, -0.2506,  0.3223,  0.3500],
[-1.1524, -0.4339, -1.2633,  0.3081],
[ 0.8487, -0.3160,  0.1198,  1.1168],
[ 0.6920, -2.1152,  1.2377, -0.2473]]])

# 这里一列不好进行view操作变为patch块，所以用个转置
y = y.transpose(2,1)
tensor([[
[-1.1258, -1.1524,  0.8487,  0.6920],
[-0.2506, -0.4339, -0.3160, -2.1152],
[ 0.3223, -1.2633,  0.1198,  1.2377],
[ 0.3500,  0.3081,  1.1168, -0.2473]]])
#然后就可以用view变成patch大小
#但用了transpose或者permute之类的函数会改变内存存储位置
#存储就变得不连续了
#所以这里用到contiguous()，不然后面有些操作就会报错
#比如要fold回来就很麻烦（意思是加了contiguous()就不麻烦了）。
y = y.view(B, self.num_patches, C, \
								 self.patch_size[0], self.patch_size[1]).contiguous()
tensor([[[[[-1.1258, -1.1524],
           [ 0.8487,  0.6920]]],


         [[[-0.2506, -0.4339],
           [-0.3160, -2.1152]]],


         [[[ 0.3223, -1.2633],
           [ 0.1198,  1.2377]]],


         [[[ 0.3500,  0.3081],
           [ 1.1168, -0.2473]]]]])
#完美结束，四个patch都被取出来了，如果有通道的话就是这样的
#x->
tensor([[[[-1.1258, -1.1524, -0.2506, -0.4339],
          [ 0.8487,  0.6920, -0.3160, -2.1152],
          [ 0.3223, -1.2633,  0.3500,  0.3081],
          [ 0.1198,  1.2377,  1.1168, -0.2473]],

         [[-1.3527, -1.6959,  0.5667,  0.7935],
          [ 0.5988, -1.5551, -0.3414,  1.8530],
          [ 0.7502, -0.5855, -0.1734,  0.1835],
          [ 1.3894,  1.5863,  0.9463, -0.8437]]]])
#unfold后
tensor([[[[[-1.1258, -1.1524],
           [ 0.8487,  0.6920]],

          [[-1.3527, -1.6959],
           [ 0.5988, -1.5551]]],


         [[[-0.2506, -0.4339],
           [-0.3160, -2.1152]],

          [[ 0.5667,  0.7935],
           [-0.3414,  1.8530]]],


         [[[ 0.3223, -1.2633],
           [ 0.1198,  1.2377]],

          [[ 0.7502, -0.5855],
           [ 1.3894,  1.5863]]],


         [[[ 0.3500,  0.3081],
           [ 1.1168, -0.2473]],

          [[-0.1734,  0.1835],
           [ 0.9463, -0.8437]]]]])

这个操作可以让每个小patch块单独进行各种操作，然后再用fold函数变为之前的维度大小，如果有需要的小伙伴可以在评论下让我更个fold的（臭不要脸