如何沿某个维度向 PyTorch 张量添加元素？

Question

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

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

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

请记住concatenation 和 stacking 之间的区别，这对于张量维度的类似问题很有帮助。

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

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