python - 如何使用 torch.stack 函数

Question

我有一个关于 torch.stack 的问题

我有 2 个张量，a.shape=(2, 3, 4) 和 b.shape=(2, 3)。 如何在没有就地操作的情况下堆叠它们？

标签： pythonpytorchtensor

Answer 1

堆叠需要相同数量的维度。一种方法是解压和堆叠。例如：

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