用数组索引火炬张量

Question

我有以下火炬张量：

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 数组索引张量的每个子元素。最好不要使用循环。

提前致谢

Answer 1

简单地说，对第一个维度使用范围(len(index))。

import torch

a = torch.tensor([[-0.2,  0.3],
    [-0.5,  0.1],
    [-0.4,  0.2]])

c = [1, 0, 1]


b = a[range(3),c]

print(b)

Answer 2

您可能想要使用 torch.gather - "Gathers values along an axis specified by 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。

用数组索引火炬张量

Index a torch tensor with an array

python

indexing

torch

pytorch

tensor