pytorch 查找指定元素的索引

在pytorch中，我们经常需要查找tensor中某一个元素的索引，可能是在向量中查找索引，也可能是在矩阵中查找索引，下面分贝举例子：

1.在矩阵中查找某个指定元素的索引：

import torch
import numpy as np

a = torch.tensor( [[1,2,3],[4,5,6],[5,6,7],[6,7,8]] )
a_t2n = a.numpy()
index = np.argwhere( a_t2n>4 )
print(index)

运行结果如下：

[[1 1]
 [1 2]
 [2 0]
 [2 1]
 [2 2]
 [3 0]
 [3 1]
 [3 2]]

返回的结果，比如第一个[1,1]代表a_t2n[1][1]=5,是满足条件的第一个元素的索引

2.在向量中查找某个指定元素的索引：

举例子如下：

import torch
import numpy as np

a = torch.tensor( [1,2,3,4,5,5,4,7] )
a_t2n = a.numpy()
index = np.argwhere( a_t2n>4 )
print(index)

运行结果如下：

[[4]
 [5]
 [7]]

第一个结果[4]，代表a_t2n[4]=5,是第一个满足这个条件的元素的索引。