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