将包含图像的张量展开为补丁

Question

我有一批尺寸为

4

、尺寸为

h x w = 180 x 320

的单通道图像。我想展开它们一系列

p

形状为

h_p x w_p

的小块，产生形状为

4 x p x h_p x w_p

的张量。如果

h

不能被

h_p

整除，或者

w

不能被

w_p

整除，则帧将以 0 填充。我尝试以下方法来实现这一目标：

import torch
tensor = torch.randn(4, 180, 320)
patch_size = (64, 64) #h_p = w_p = 64
unfold = torch.nn.Unfold(kernel_size=patch_size, stride=patch_size, padding=0)
unfolded = unfold(tensor)
print(unfolded.shape)

打印：

torch.Size([16384, 10])

我在这里缺少什么？

Answer 1

我输入了形状

[#batches, height, width] = [4,180,320]

。我想展开它们一系列

p

形状为

h_p x w_p

的小块，产生形状为

4 x p x h_p x w_p

的张量。请注意，要使用大小为

h x w = 180 x 320

的补丁覆盖所有

h_p x w_p = 64 x 64

元素，我将需要

p = 3 x 5 = 15

补丁：

所以，我在两侧添加了 6 的 padding。其余代码我已在评论中解释：

patch_size = (64,64)
input = torch.randn(4,180,320)

# adding of 6 on top and bottom, to make up total padding of 12 rows, 
# so that our frame will become of size 192 x 320 and we can fit 3
# kernels of size 64 x 64 vertically
input = f.pad(input, pad=(0,0,6,6))
print(input.shape) # [4,192,320]

# add additional dimension indicating single channel
input = input.unsqueeze(1) # [4,1,192, 320]
print(input.shape)

# unfold with both stride and kernel size of 64 x 64
unfold = torch.nn.Unfold(kernel_size=patch_size, stride=(64,64))
unfolded = unfold(input)
print(unfolded.shape) # [4, 4096, 15] 
# 4 for batch size
# 4096 = 64 x 64 elements in one patch
# 15 = we can fit 15 patches of size 64 x 64 in frame of size 192 x 329

# reshape result to desired size
# size(0) = 4 = batch size
# -1 to infer p or number of patches, by our calculations it will be 15
# *patch_size = 64 x 64
unfolded = unfolded.view(unfolded.size(0),-1,*patch_size) 
print(unfolded.shape) # [4, 15, 64, 64]

正确输出：

torch.Size([4, 192, 320])
torch.Size([4, 1, 192, 320])
torch.Size([4, 4096, 15])
torch.Size([4, 15, 64, 64]