Pytorch 如何像 for 循环一样堆叠张量

Question

我想连接在for循环中生成的张量，并得到2dTensor。
标准 Python，如下所示。

li = []
for i in range(0, len(items)):
# calc something
li.append(calc_result)

就我而言，在 for 循环中，生成torch.Size([768])张量，我想得到torch.Size([len(item),768])张量。
怎么做？

Answer 1

您可以使用torch.stack：

torch.stack(li, dim=0)

之后 for 循环将为您提供该大小的torch.Tensor。

请注意，如果您事先知道最终张量的大小，则可以事先分配一个空张量并将其填充到 for 循环中：

x = torch.empty(size=(len(items), 768))
for i in range(len(items)):
x[i] = calc_result

这通常比执行堆栈更快。

Answer 2

使用torch.stack的接受答案是不正确的，因为它插入了一个额外的维度，给出了形状[1, len(items), 768]的张量。

请改用torch.vstack：

torch.vstack(li, dim=0)

得到形状[len(items), 768]的张量。