我有以下火炬张量:
tensor([[-0.2, 0.3],
[-0.5, 0.1],
[-0.4, 0.2]])
以及以下 numpy 数组:(如有必要,我可以将其转换为其他内容)
[1 0 1]
我想获得以下张量:
tensor([0.3, -0.5, 0.2])
即我希望 numpy 数组索引张量的每个子元素。最好不使用循环。
提前致谢
最佳答案
您可能想使用 torch.gather
- “沿由dim 指定的轴收集值。”
t = torch.tensor([[-0.2, 0.3],
[-0.5, 0.1],
[-0.4, 0.2]])
idxs = np.array([1,0,1])
idxs = torch.from_numpy(idxs).long().unsqueeze(1)
# or torch.from_numpy(idxs).long().view(-1,1)
t.gather(1, idxs)
tensor([[ 0.3000],
[-0.5000],
[ 0.2000]])
在这里,您的索引是 numpy 数组,因此您必须将其转换为 LongTensor。
关于python - 用数组索引 torch 张量,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/61311688/