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

Question

在pytorch中，给定形状(1X11)的张量a和形状(1X11)的b，torch.stack((a,b),0)会给我一个形状的张量(2X11)

但是，当 a 的形状为 (2X11) 且 b 的形状为 (1X11) 时，torch.stack((a,b),0) 将引发错误 cf。 "the two tensor size must exactly be the same".

因为这两个张量是一个模型的输出（包含梯度），我无法将它们转换为 numpy 以使用 np.stack() 或 np.vstack().

是否有最小 GPU 内存使用率的解决方案？

Answer 1

您似乎想使用 torch.cat() (concatenate tensors along an existing dimension) and not torch.stack()（concatenate/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])