Python 如何使用torch.stack函数

我有一个关于火炬的问题

我有两个张量，a.shape=（2,3,4）和b.shape=（2,3）。

---

如何在不进行就地操作的情况下堆叠它们？

堆叠需要相同数量的尺寸。一种方法是解列并堆叠。例如：

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

>>> 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]]))

---

使用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])

下面是另一个例子：

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

>>> 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

参数指定在哪个维度上堆叠维度相等的张量。

假设有两个维度相等的张量a，b，即a（a，b，C）so 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

---

如果他们不一样的话，它就不会起作用。小心

您想要的是使用

unsqueze

，就像您所做的那样。创建一个新维度，并且提供的所有张量必须相同大小。对于@drevicko正确提到的

torch.stack（[a，b]，dim=2）

，此答案不正确

torch.cat

连接给定维度中的序列，而

torch.stack

连接新维度中的序列，如下所述：。此操作不会运行。相反，您将收到“RuntimeError:stack期望每个张量大小相等，但在条目0处得到[2,3,4]，在条目1处得到[2,3,1]”