Pytorch张量获取具有特定值的元素的索引？

Question

我有两个张量，张量a和张量b

我想得到张量b中所有值的索引

比如说

a = torch.Tensor([1,2,2,3,4,4,4,5])

b = torch.Tensor([1,2,4])

我想要张量a中的1, 2, 4的索引。我可以通过以下代码来实现这一点

a = torch.Tensor([1,2,2,3,4,4,4,5])
b = torch.Tensor([1,2,4])
mask = torch.zeros(a.shape).type(torch.bool)
print(mask)
for e in b:
    mask = mask + (a == e)
    print(mask)

没有for我怎么做

Answer 1

如果您不想使用for循环，可以使用列表理解：

mask = [a[index] for index in b]

如果你甚至不想使用“for”这个词，你可以将张量转换成numpy，并使用numpy索引

mask = torch.tensor(a.numpy()[b.numpy()])

更新

可能误解了你的问题。在这种情况下，我想说实现这一点的最佳方法是通过列表理解。（切片可能无法实现这一点

mask = [index for index,value in enumerate(a) if value in b.tolist()] 

它迭代a中的每个元素，获取它们的索引和值，如果值在b中，则获取索引