pytorch按照指定元素索引排序选择操作大全（nonzero、sort、index_select函数）

一、特定元素的位置索引

获取矩阵中特定元素的位置索引，可以通过nonzero()结合判断条件来获取。nonzero()函数可以返回非0元素的位置。下面举一下一维以及二维的相关例子。

（1）一维矩阵的nonzero()

>>> import torch
>>> a = torch.tensor([1,0,2,30,0,1])
>>> a.nonzero()
tensor([[0],
        [2],
        [3],
        [5]])
>>> a.nonzero().size()
torch.Size([4, 1])

（2）二维矩阵的nonzero()

>>> b = torch.tensor([[1,0,2],[0,2,1]])
>>> b
tensor([[1, 0, 2],
        [0, 2, 1]])
>>> b.nonzero()
tensor([[0, 0],
        [0, 2],
        [1, 1],
        [1, 2]])
>>> b.nonzero().size()
torch.Size([4, 2])

（3）结合判断条件的nonzero()，可以获取指定元素的位置索引

下面例子是获取二维矩阵c中2对应的索引位置。

>>> c = torch.tensor([[1,0,2],[0,2,1]]) 
>>> c == torch.tensor(2)
tensor([[False, False,  True],
        [False,  True, False]])
>>> (c == torch.tensor(2)).nonzero()
tensor([[0, 2],
        [1, 1]])
>>> (c == torch.tensor(2)).nonzero().size()
torch.Size([2, 2])

（4）torch.nonzero()函数

torch.nonzero(..., as_tuple=False/True)

as_tuple(默认是False)，是否作为元组返回每一维的张量位置

• if false：输出的每一行为非零元素的索引
• if true：输出是每一个维度都有一个一维的张量

>>> a = torch.tensor([[1,2,0],[0,2,6]])
>>> torch.nonzero(a)
tensor([[0, 0],
        [0, 1],
        [1, 1],
        [1, 2]])
>>> torch.nonzero(a,as_tuple=False)
tensor([[0, 0],
        [0, 1],
        [1, 1],
        [1, 2]])
>>> torch.nonzero(a,as_tuple=True)
(tensor([0, 0, 1, 1]), tensor([0, 1, 1, 2]))

二、torch.sort()排序操作

>>> d = torch.tensor([0.1215, 0.1256, 0.1227, 0.1343])
>>> sorted, indices = torch.sort(d, descending=True)  #descending：递减（从大到小）
>>> sorted
tensor([0.1343, 0.1256, 0.1227, 0.1215])
>>> indices
tensor([3, 1, 2, 0])
>>> sorted, indices = torch.sort(d, descending=False)
>>> sorted
tensor([0.1215, 0.1227, 0.1256, 0.1343])
>>> indices
tensor([0, 2, 1, 3])

三、torch.index_select()根据索引位置序列进行选择拼接

用法：

torch.index_select(input, dim, index, *, out=None) → Tensor

• input (Tensor) – the input tensor.
• dim (int) – the dimension in which we index
• index (IntTensor or LongTensor) – the 1-D tensor containing the indices to index

第一个参数是索引的对象，第二个参数0表示按行索引，1表示按列进行索引，第三个参数是一个tensor，就是索引序列的序号。

>>> a = torch.linspace(1, 12, steps=12).view(3, 4)
>>> a
tensor([[ 1.,  2.,  3.,  4.],
        [ 5.,  6.,  7.,  8.],
        [ 9., 10., 11., 12.]])
>>> a.index_select(0, torch.tensor([0, 2]))
tensor([[ 1.,  2.,  3.,  4.],
        [ 9., 10., 11., 12.]])
>>> torch.index_select(a, 0, torch.tensor([0, 2]))
tensor([[ 1.,  2.,  3.,  4.],
        [ 9., 10., 11., 12.]])
>>> torch.index_select(a, 1, torch.tensor([1, 3]))
tensor([[ 2.,  4.],
        [ 6.,  8.],
        [10., 12.]])

四、结合上面函数使用的实际例子

目的：是将按照指定的元素（类别）进行筛选以及排序，选出你想获取的特征或者数值。（根据index里的指定元素，获取另外对应矩阵对应的数值。）

# cls_prob = F.softmax(cls_score, 1)                     # 将预测的结果转换为各个类别的0-1的概率
# confidence, index = torch.max(tgt_cls_prob, 1)         # 选出对应最大概率的置信度以及类别

#### eg：
import torch

confidence = torch.tensor([0.1215, 0.1256, 0.1227, 0.1343, 0.1157, 0.1139, 0.1086, 0.0927, 0.0805,
        0.0896, 0.0667, 0.0782, 0.0984, 0.0657, 0.0818, 0.1009, 0.0966, 0.0639,
        0.0860, 0.0703, 0.1160, 0.0897, 0.0728, 0.0785, 0.0811, 0.1490, 0.0837,
        0.0805, 0.1205, 0.0930, 0.1165, 0.0828, 0.0684, 0.0796, 0.1050, 0.0982,
        0.0712, 0.1210, 0.3321, 0.1182, 0.0978, 0.1038, 0.0995, 0.0748, 0.1117,
        0.0736, 0.0825, 0.0828, 0.0930, 0.0528, 0.0893, 0.0636, 0.1043, 0.0848,
        0.0820, 0.0975, 0.1010, 0.0925, 0.0716, 0.0746, 0.0808, 0.0750, 0.0877,
        0.0668, 0.0840, 0.0708, 0.0732, 0.1356, 0.0715, 0.0932, 0.0997, 0.0870,
        0.0778, 0.0987, 0.0744, 0.0709, 0.0977, 0.0708, 0.0972, 0.0996, 0.0786,
        0.0750, 0.1059, 0.0712, 0.0663, 0.0993, 0.0726, 0.0986, 0.0772, 0.0803,
        0.0702, 0.1130, 0.1043, 0.0686, 0.0927, 0.0681, 0.1205, 0.0810, 0.0932,
        0.0818, 0.0852, 0.1034, 0.0765, 0.0693, 0.0857, 0.0771, 0.0871, 0.0870,
        0.0932, 0.0952, 0.0701, 0.0803, 0.0996, 0.0800, 0.0854, 0.0932, 0.0795,
        0.0826, 0.1255, 0.0712, 0.1643, 0.0535, 0.0783, 0.0650, 0.0748, 0.0757,
        0.0751, 0.0788])

index = torch.tensor([ 3,  3,  3,  3,  3,  3,  3, 19, 19,  4, 11,  3,  3,  3,  3,  3,  3,  3,
         8,  3,  3,  3,  3, 11,  3,  4, 11,  8,  3,  3,  4,  3,  5,  0, 11,  3,
        11, 11, 19,  3,  0,  7, 11,  1,  3, 11,  8, 11, 19, 16, 11,  1, 11,  3,
         3,  3,  4,  3,  8,  8,  3,  3, 11,  3,  2,  8,  3,  4, 20,  3,  3, 11,
        18, 11,  4, 11,  8,  3,  3,  4, 11,  3,  4,  3,  6, 11,  8,  3,  3, 19,
         3,  3, 11,  3,  8,  3,  3,  3, 19,  8,  3,  3,  3,  1,  3,  8, 19,  3,
         3, 11, 11,  1, 11, 11, 11, 19,  8, 18,  3, 11, 11, 16,  8,  8, 11,  8,
         3, 18])

cls_index = (index == torch.tensor(8)).nonzero()  
print(cls_index)
# tensor([[ 18],
#         [ 27],
#         [ 46],
#         [ 58],
#         [ 59],
#         [ 65],
#         [ 76],
#         [ 86],
#         [ 94],
#         [ 99],
#         [105],
#         [116],
#         [122],
#         [123],
#         [125]])
print(cls_index.size()) # torch.Size([15, 1])
cls_index = (index == torch.tensor(8)).nonzero().squeeze(1) 
print(cls_index)        # 类别在原数组(confidence)的位置索引
# tensor([ 18,  27,  46,  58,  59,  65,  76,  86,  94,  99, 105, 116, 122, 123, 125])
print(cls_index.size()) # torch.Size([15])
print(confidence[cls_index])   # 获取该元素(类别)对应置信度(confidence)的数值
# tensor([0.0860, 0.0805, 0.0825, 0.0716, 0.0746, 0.0708, 0.0977, 0.0726, 0.0927,
#         0.0818, 0.0771, 0.0795, 0.0783, 0.0650, 0.0757])
sorted, indices = torch.sort(confidence[cls_index], descending=True) # 降序  (针对该类别的置信度进行排序，此时sorted和indices是对应)
print(sorted)   # 当前类别局部置信度排序
# tensor([0.0977, 0.0927, 0.0860, 0.0825, 0.0818, 0.0805, 0.0795, 0.0783, 0.0771,
#        0.0757, 0.0746, 0.0726, 0.0716, 0.0708, 0.0650])
print(indices)  # 当前类别局部索引排序
# tensor([ 6,  8,  0,  2,  9,  1, 11, 12, 10, 14,  4,  7,  3,  5, 13])
print(indices[0:10])   # 当前类别局部排序top 10
# tensor([ 6,  8,  0,  2,  9,  1, 11, 12, 10, 14])
print(cls_index[indices[0:10]])  # 对应原数组(confidence)该类别对应最高置信度top10的索引
# tensor([ 76,  94,  18,  46,  99,  27, 116, 122, 105, 125])

confidences = torch.index_select(confidence,0,cls_index[indices[0:10]]) 
print(confidences)
# tensor([0.0977, 0.0927, 0.0860, 0.0825, 0.0818, 0.0805, 0.0795, 0.0783, 0.0771, 0.0757])
# index_select里面的confidence一般可以换为对应的特征向量，然后就是将特征向量按照指定位置或通道进行拼接之类的
print(confidences.size())
# torch.Size([10])

参考：
index_select()的用法
pytorch官方使用文档