我该如何使用torch.stack？

Question

如何使用torch.stack在不进行原地操作的情况下堆叠两个形状分别为a.shape = (2, 3, 4)和b.shape = (2, 3)的张量？

- Suho Cho

Answer 1

堆叠需要相同数量的维度。一种方法是先使用unsqueeze增加维度，再进行堆叠。例如：

a.size()  # 2, 3, 4
b.size()  # 2, 3
b = torch.unsqueeze(b, dim=2)  # 2, 3, 1
# torch.unsqueeze(b, dim=-1) does the same thing

torch.stack([a, b], dim=2)  # 2, 3, 5

- arjoonn

---

你希望使用torch.cat和unsqueeze，就像你已经做过的那样。torch.stack会创建一个新的维度，所有提供的张量都必须具有相同的大小。 - drevicko

这个答案使用了 torch.stack([a, b], dim=2) 是错误的，正确的方法是使用 torch.cat([a,b], dim=2)，就像 @drevicko 所提到的那样。torch.cat 在给定维度上连接序列，而 torch.stack 则在一个新维度上堆叠序列，详情请见：https://dev59.com/EVQJ5IYBdhLWcg3wNjO4#54307331 。 - warriorUSP

无法运行，相反你会收到'RuntimeError: stack希望每个张量具有相同的尺寸，但是在entry 0有[2,3,4]而在entry 1有[2,3,1]'的错误提示信息。 - JP Zhang

Answer 2

在使用PyTorch 1.2或1.4时，arjoonn的回答对我没有用。

我使用了PyTorch 1.2和1.4中的torch.cat，而不是torch.stack：

>>> import torch
>>> a = torch.randn([2, 3, 4])
>>> b = torch.randn([2, 3])
>>> b = b.unsqueeze(dim=2)
>>> b.shape
torch.Size([2, 3, 1])
>>> torch.cat([a, b], dim=2).shape
torch.Size([2, 3, 5])

如果您想使用torch.stack，则张量的维度必须相同：

>>> a = torch.randn([2, 3, 4])
>>> b = torch.randn([2, 3, 4])
>>> torch.stack([a, b]).shape
torch.Size([2, 2, 3, 4])

这里有另一个例子：

>>> t = torch.tensor([1, 1, 2])
>>> stacked = torch.stack([t, t, t], dim=0)
>>> t.shape, stacked.shape, stacked

(torch.Size([3]),
 torch.Size([3, 3]),
 tensor([[1, 1, 2],
         [1, 1, 2],
         [1, 1, 2]]))

使用stack函数时，可以通过dim参数指定在哪个维度上堆叠具有相同维度的张量。

- gil.fernandes

Answer 3

假设您有两个张量a、b，它们在维度上相等，即a（A，B，C），b（A，B，C）。

一个例子：

a=torch.randn(2,3,4)
b=torch.randn(2,3,4)
print(a.size())  # 2, 3, 4
print(b.size()) # 2, 3, 4

f=torch.stack([a, b], dim=2)  # 2, 3, 2, 4
f

如果它们的维度不同，IT技术就不会起作用。要小心！！

- pourya