numpy实现torch的topk方法

torch中提供了topk方法用来返回矩阵中对应维度中最大的K个元素以及在对应维度中的index，但是numpy并没有提供和torch一样的topk方法，所以在这里通过numpy的argpartition实现torch中的topk方法。

直接给出代码：

def topk_(matrix, K, axis=1):
    if axis == 0:
        row_index = np.arange(matrix.shape[1 - axis])
        topk_index = np.argpartition(-matrix, K, axis=axis)[0:K, :]
        topk_data = matrix[topk_index, row_index]
        topk_index_sort = np.argsort(-topk_data,axis=axis)
        topk_data_sort = topk_data[topk_index_sort,row_index]
        topk_index_sort = topk_index[0:K,:][topk_index_sort,row_index]
    else:
        column_index = np.arange(matrix.shape[1 - axis])[:, None]
        topk_index = np.argpartition(-matrix, K, axis=axis)[:, 0:K]
        topk_data = matrix[column_index, topk_index]
        topk_index_sort = np.argsort(-topk_data, axis=axis)
        topk_data_sort = topk_data[column_index, topk_index_sort]
        topk_index_sort = topk_index[:,0:K][column_index,topk_index_sort]
    return topk_data_sort, topk_index_sort

测试功能：

# torch.topk方法
>>> a=torch.rand(5,4)
>>> a
tensor([[0.0154, 0.5266, 0.6294, 0.6897],
        [0.2201, 0.7039, 0.2639, 0.0681],
        [0.1006, 0.0464, 0.3314, 0.2052],
        [0.3954, 0.4373, 0.2147, 0.3532],
        [0.6560, 0.0549, 0.8040, 0.3528]])
>>> a.topk(2,dim=1,largest=True,sorted=True)
values=tensor([[0.6897, 0.6294],
        [0.7039, 0.2639],
        [0.3314, 0.2052],
        [0.4373, 0.3954],
        [0.8040, 0.6560]]),
indices=tensor([[3, 2],
        [1, 2],
        [2, 3],
        [1, 0],
        [2, 0]]))
>>> a.topk(2,dim=0,largest=True,sorted=True)
torch.return_types.topk(
values=tensor([[0.6560, 0.7039, 0.8040, 0.6897],
        [0.3954, 0.5266, 0.6294, 0.3532]]),
indices=tensor([[4, 1, 4, 0],
        [3, 0, 0, 3]]))

# 自定义numpy的topk方法
>>> a_np=np.array(a)
>>> a_np
array([[0.01537341, 0.5266498 , 0.6293524 , 0.689658  ],
       [0.2201249 , 0.70394784, 0.26386315, 0.06814277],
       [0.10058308, 0.04639381, 0.3313678 , 0.20519769],
       [0.395352  , 0.43731135, 0.21468669, 0.35324287],
       [0.655955  , 0.05492574, 0.80404747, 0.35280174]], dtype=float32)
>>> topk(a_np,2,axis=1)
(array([[0.689658  , 0.6293524 ],
       [0.70394784, 0.26386315],
       [0.3313678 , 0.20519769],
       [0.43731135, 0.395352  ],
       [0.80404747, 0.655955  ]], dtype=float32), array([[3, 2],
       [1, 2],
       [2, 3],
       [1, 0],
       [2, 0]]))
>>> topk(a_np,2,axis=0)
(array([[0.655955  , 0.70394784, 0.80404747, 0.689658  ],
       [0.395352  , 0.5266498 , 0.6293524 , 0.35324287]], dtype=float32), array([[4, 1, 4, 0],
       [3, 0, 0, 3]]))

可以发现已经完全实现了和torch.topk相同的功能～