PyTorch使用torch.sort()函数来筛选出前k个最大的项或者筛选出前k个最小的项

参考链接: torch.topk(input, k, dim=None, largest=True, sorted=True, out=None) -＞ (Tensor, LongTensor)

筛选出前k个最小的数:

Microsoft Windows [版本 10.0.18363.1256]
(c) 2019 Microsoft Corporation。保留所有权利。

C:\Users\chenxuqi>conda activate ssd4pytorch1_2_0

(ssd4pytorch1_2_0) C:\Users\chenxuqi>python
Python 3.7.7 (default, May  6 2020, 11:45:54) [MSC v.1916 64 bit (AMD64)] :: Anaconda, Inc. on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.manual_seed(seed=20200910)
<torch._C.Generator object at 0x0000026DB7CFD330>
>>>
>>> data = torch.randint(100,(15,))
>>> data
tensor([63, 48, 14, 47, 28,  5, 80, 68, 88, 61,  6, 84, 82, 87, 59])
>>> # 筛选出前k个最小的数
>>> k = 7
>>> a, idx1 = torch.sort(data)
>>> b, idx2 = torch.sort(idx1)
>>> a
tensor([ 5,  6, 14, 28, 47, 48, 59, 61, 63, 68, 80, 82, 84, 87, 88])
>>> data
tensor([63, 48, 14, 47, 28,  5, 80, 68, 88, 61,  6, 84, 82, 87, 59])
>>> data[idx2<k]
tensor([48, 14, 47, 28,  5,  6, 59])
>>>
>>>
>>>

筛选出前k个最大的数:

Microsoft Windows [版本 10.0.18363.1256]
(c) 2019 Microsoft Corporation。保留所有权利。

C:\Users\chenxuqi>conda activate ssd4pytorch1_2_0

(ssd4pytorch1_2_0) C:\Users\chenxuqi>python
Python 3.7.7 (default, May  6 2020, 11:45:54) [MSC v.1916 64 bit (AMD64)] :: Anaconda, Inc. on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.manual_seed(seed=20200910)
<torch._C.Generator object at 0x000001846D60D330>
>>>
>>>
>>> data = torch.randint(100,(15,))
>>> data
tensor([63, 48, 14, 47, 28,  5, 80, 68, 88, 61,  6, 84, 82, 87, 59])
>>> # 筛选出前k个最大的数
>>>
>>> k = 7
>>> a, idx1 = torch.sort(data, descending=True)
>>> b, idx2 = torch.sort(idx1)
>>>
>>> a
tensor([88, 87, 84, 82, 80, 68, 63, 61, 59, 48, 47, 28, 14,  6,  5])
>>> data
tensor([63, 48, 14, 47, 28,  5, 80, 68, 88, 61,  6, 84, 82, 87, 59])
>>> data[idx2<k]
tensor([63, 80, 68, 88, 84, 82, 87])
>>>
>>>
>>>

简单证明:

反向传播,只传播到前K个最大的项:

Microsoft Windows [版本 10.0.18363.1316]
(c) 2019 Microsoft Corporation。保留所有权利。

C:\Users\chenxuqi>conda activate ssd4pytorch1_2_0

(ssd4pytorch1_2_0) C:\Users\chenxuqi>python
Python 3.7.7 (default, May  6 2020, 11:45:54) [MSC v.1916 64 bit (AMD64)] :: Anaconda, Inc. on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.manual_seed(seed=20200910)
<torch._C.Generator object at 0x000001F7F101D330>
>>>
>>> data = torch.randn(15,requires_grad=True)
>>> data
tensor([ 0.2824, -0.3715,  0.9088, -1.7601, -0.1806,  2.0937,  1.0406, -1.7651,
         1.1216,  0.8440,  0.1783,  0.6859, -1.5942, -0.2006, -0.4050],
       requires_grad=True)
>>>
>>> # 筛选出前k个最大的数
>>> k = 7
>>> a, idx1 = torch.sort(data, descending=True)
>>> b, idx2 = torch.sort(idx1)
>>> a
tensor([ 2.0937,  1.1216,  1.0406,  0.9088,  0.8440,  0.6859,  0.2824,  0.1783,
        -0.1806, -0.2006, -0.3715, -0.4050, -1.5942, -1.7601, -1.7651],
       grad_fn=<SortBackward>)
>>> b
tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14])
>>> idx1
tensor([ 5,  8,  6,  2,  9, 11,  0, 10,  4, 13,  1, 14, 12,  3,  7])
>>> idx2
tensor([ 6, 10,  3, 13,  8,  0,  2, 14,  1,  4,  7,  5, 12,  9, 11])
>>>
>>> a
tensor([ 2.0937,  1.1216,  1.0406,  0.9088,  0.8440,  0.6859,  0.2824,  0.1783,
        -0.1806, -0.2006, -0.3715, -0.4050, -1.5942, -1.7601, -1.7651],
       grad_fn=<SortBackward>)
>>> data
tensor([ 0.2824, -0.3715,  0.9088, -1.7601, -0.1806,  2.0937,  1.0406, -1.7651,
         1.1216,  0.8440,  0.1783,  0.6859, -1.5942, -0.2006, -0.4050],
       requires_grad=True)
>>> data[idx2<k]
tensor([0.2824, 0.9088, 2.0937, 1.0406, 1.1216, 0.8440, 0.6859],
       grad_fn=<IndexBackward>)
>>> sum_topK = data[idx2<k].sum()
>>> sum_topK
tensor(6.9770, grad_fn=<SumBackward0>)
>>>
>>> 2.0937+1.1216+1.0406+0.9088+0.8440+0.6859+0.2824
6.977000000000001
>>> data.grad
>>> print(data.grad)
None
>>> sum_topK.backward()
>>> print(data.grad)
tensor([1., 0., 1., 0., 0., 1., 1., 0., 1., 1., 0., 1., 0., 0., 0.])
>>>
>>>
>>>