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

numpy实现torch的topk方法相关推荐

numpy和torch数据操作对比
对numpy和torch数据操作进行对比,避免遗忘. ndarray和tensor import torch import numpy as npnp_data = np.arange(6).resh ...
list,numpy,tensor之间相互转换的方法
list,numpy,tensor之间相互转换的方法: a=[[1,2],[3,4]]#list print(a) b=np.array(a)#list->numpy print(b) c=to ...
Pytorch函数之topk()方法
根据Pytorch中的手册可以看到,topk()方法用于返回输入数据中特定维度上的前k个最大的元素. torch.topk(input, k, dim=None, largest=True, sort ...
pytorch下Numpy，Torch，Spicy，NetworkX及其他基本数据类型相关操作(持续更新)
Tricks 1. torch.sparse.FloatTensor(position, value) 稀疏张量表示为一对稠密张量:一个值张量和一个二维指标张量(每一维中存储多个值).一个稀疏张量可以 ...
PyTorch中的topk方法以及分类Top-K准确率的实现
PyTorch中的topk方法以及分类Top-K准确率的实现 Top-K 准确率在分类任务中的类别数很多时(如ImageNet中1000类),通常任务是比较困难的,有时模型虽然不能准确地将groun ...
python增加一列数据_Python编程给numpy矩阵添加一列方法示例
首先我们有一个数据是一个mn的numpy矩阵现在我们希望能够进行给他加上一列变成一个m(n+1)的矩阵 import numpy as np a = np.array([[1,2,3],[4,5,6] ...
PyTorch | 通过torch.normal()创建概率分布的张量 | torch.normal()如何使用？torch.normal()使用方法 | torch.normal()例子
在公众号[计算机视觉联盟]后台回复[9076]获取独家200页手推AI笔记:我的微信:PursueWin: --by Sophia 中科院学霸 | 上市AI算法工程师 | CSDN博客专家通过 ...
PyTorch | 通过torch.eye创建单位对角矩阵 | torch.eye()如何使用？torch.eye()例子 | torch.eye()使用方法
在公众号[计算机视觉联盟]后台回复[9076]获取独家200页手推AI笔记:我的微信:PursueWin: --by Sophia 中科院学霸 | 上市AI算法工程师 | CSDN博客专家 t ...
PyTorch | torch.linspace()创建均分数列张量 | torch.linspace()如何使用？| torch.linspace()使用方法 | torch.linspace例子
公众号[计算机视觉联盟]后台回复[PyTorch]可以获得独家PyTorch学习教程pdf版通过torch.linspace创建均分数列张量步长=(Start - end)/(Steps - 1 ...

numpy实现torch的topk方法

numpy实现torch的topk方法相关推荐

最新文章

热门文章