关于 pytorch 改变张量形状的问题

Question

当我想将图像分割成多个补丁时，我首先想到的是使用 pytorch view() function。 例如，形状为 (1,3,256,256)（pytorch 样式）的图像，并将其拆分为 8x8=64 个块，每个块的高度和宽度为 32。因此对于此图像，我们可以得到 256/32=8 的行和8 列补丁，所以我们总共有 8x8=64 个补丁。

我想将图像（1,3,256,256）分割成补丁，每个补丁的形状是（1,3,32,32），并将这些张量重塑为（1,8x8,32x32x3）的形状，这里8x8是补丁的数量, 32x32x3 是补丁高度 * 宽度 * 通道。

下面的代码使用 einops 重新排列 function 可以获得正确答案，但是当我使用视图 function 时，我得到了正确的形状但不正确的值张量。 谁能告诉我如何通过视图function来实现这个操作？

from einops.layers.torch import Rearrange
img = torch.randn(1, 3, 256, 256)
import copy
img2 = copy.deepcopy(img)
b, c, h, w = img.size()
p=32
to_patch_embedding = nn.Sequential(
    Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=32, p2=32),
)
img2 = img2.view(b, h // p * w // p, c * p * p)


print(img2.shape)
print(img2==to_patch_embedding(img))

----------------------------------------output--------------------------------------------
torch.Size([1, 64, 3072])
tensor([[[ True, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     ...,
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False, False],
     [False, False, False,  ..., False, False,  True]]])

Answer 1

您可以按照pytorch 讨论中提到的方式尝试这种方式。

import torch
a = torch.randn(1, 3, 256, 256)
a = a.unfold(2, 32, 32).unfold(3, 32, 32)
a = a.contiguous().view(a.size(0), a.size(2)*a.size(3), a.size(1)*a.size(-1)*a.size(-2))
print(a.shape)

output 形状如您所料：

torch.Size([1, 64, 3072])

希望这对你有用。