如何使用pytorch将二维张量分成更小的块？

Question

我已经下载了 EMNIST 字母数据集，并将每个图像转换为形状为

torch.tensor

的

torch.size([28, 28])

对象。但是，我想将 28*28 图像分成 7*7 块，每个块大小为 16。

即如果图像像素从左到右、从上到下标记为 1, 2, ..., 784

[
  [1, 2, 3, ..., 28],
  ...
  [        ..., 784]
]

我希望输出是一个

torch.tensor

大小为

torch.size([7, 7, 16])

的对象

[
  [
    [1, 2, 3, 4, 29, 30, 31, 32, 57, 58, 59, 60, 85, 86, 87, 88],
    ...
    [25, 26, 27, 28, 53, 54, 55, 56, 81, 82, 83, 84, 109, 110, 111, 112]
  ],
  ...
  [
    ...
    [697, 698, 699, 700, 725, 726, 727, 728, 753, 754, 755, 756, 781, 782, 783, 784]
  ]
]

我尝试过使用

torch.view(7, 7, 16)

，但它没有显示出预期的结果。

非常感谢^_^

Answer 1

使用

torch.nn.functional.unfold

。它速度更快，如果需要的话可以微分，你甚至可以获取重叠的补丁。这是一个例子：

x = torch.arange(1,28*28+1).view(28,28).float() # unfold only works with float tensor
out = torch.nn.functional.unfold(x.unsqueeze(0).unsqueeze(0), kernel_size= 4, dilation= 1, padding= 0, stride= 4) # 2 unsqueeze to make `x` have dim 4 (BxCxHxW) 
out.permute(0,2,1).shape  # torch.Size([1, 49, 16])

输出