Pytorch - 堆栈维度必须完全相同？

Question

在pytorch,鉴于张量a的形状(1X11)和 b造型(1X11),torch.stack((a,b),0)都会给我形状的张量(2X11)

但是,当a形状(2X11)和b形状时(1X11),torch.stack((a,b),0)会引起错误cf. "两个张量大小必须完全相同".

因为两个张量是模型的输出(包括渐变),我无法将它们转换为numpy使用np.stack()或np.vstack().

是否有任何可能的解决方案,至少GPU内存使用？

Answer 1

您似乎想要使用torch.cat()(沿现有维度连接张量)而不是torch.stack()(沿新维度连接/堆栈张量):

import torch

a = torch.randn(1, 42, 1, 1)
b = torch.randn(1, 42, 1, 1)

ab = torch.stack((a, b), 0)
print(ab.shape)
# torch.Size([2, 1, 42, 1, 1])

ab = torch.cat((a, b), 0)
print(ab.shape)
# torch.Size([2, 42, 1, 1])
aab = torch.cat((a, ab), 0)
print(aab.shape)
# torch.Size([3, 42, 1, 1])