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相同的功能~