如何沿某个维度将元素添加到 PyTorch 张量？答案

【问题标题】：How can I add an element to a PyTorch tensor along a certain dimension?如何沿某个维度将元素添加到 PyTorch 张量？
【发布时间】：2020-07-20 22:29:55
【问题描述】：

我有一个张量inps，其大小为[64, 161, 1]，我有一些新数据d，其大小为[64, 161]。如何将d 添加到inps 以使新大小为[64, 161, 2]？

【问题讨论】：

附带说明，尽管这是一个非常常见的问题，但我找不到实际上以完全相同的方式解决问题的答案。如果你能找到一个重复的问题，请随意链接。

标签： python pytorch tensor

【解决方案1】：

使用.unsqueeze() 和torch.cat() 有一种更简洁的方式，直接使用PyTorch 接口：

import torch

# create two sample vectors
inps = torch.randn([64, 161, 1])
d = torch.randn([64, 161])

# bring d into the same format, and then concatenate tensors
new_inps = torch.cat((inps, d.unsqueeze(2)), dim=-1)
print(new_inps.shape)  # [64, 161, 2]

本质上，解压缩第二维已经使两个张量具有相同的形状；你只需要小心沿着正确的维度展开。同样，不幸的是，concatenation 的命名与其他类似命名的 NumPy 函数不同，但行为相同。请注意，不是让torch.cat 通过提供dim=-1 来计算维度，您还可以显式提供要连接的维度，在这种情况下，将其替换为dim=2。

记住difference between concatenation and stacking，这对于张量维度的类似问题很有帮助。

【讨论】：

如果您知道它始终是最后一个维度，将 2 更改为 -1 会使其更通用。另外，我不知道为什么，我更喜欢d[..., None] 而不是使用unsqueeze :)
好收获！我自己很少使用它（主要是为了明确改变什么维度，但为了概括起见，我将编辑答案。

【解决方案2】：

您必须首先重塑d，使其具有第三个维度，沿着该维度可以进行连接。在它具有第三维并且两个张量具有相同的维数之后，您可以使用 torch.cat((inps, d),2) 将它们堆叠起来。

old_shape = tuple(d.shape)
new_shape = old_shape + (1,)
inps_new = torch.cat( (inps, d.view( new_shape ), 2)

【讨论】：

我也可以使用unsqueeze吗？

【解决方案3】：

或者，您可以通过挤压更大的张量和堆叠来实现这一点：

inps = torch.randn([64, 161, 1])
d = torch.randn([64, 161])

res = torch.stack((inps.squeeze(), d), dim=-1)

res.shape
>>> [64, 161, 2]

【讨论】：