以二维numpy矩阵为例
import torch
import numpy as np
K=3 #取每行最小3个值的索引
data=np.random.rand(4,7)
print(data)
data=torch.from_numpy(data)
a, idx = torch.sort(data, descending=False)
lists=idx[:,:K]
print(lists)
运行结果如下：
results.jpg
import torch
import numpy as np
K=3 #取每行最小3个值的索引
data=np.random.rand(4,7)
print(data)
data=torch.from_numpy(data)
a, idx = torch.sort(data, descending=False)
lists=idx[:,:K]
print(lists)
本文标题：pytorch 计算2维矩阵每行最小N个值的索引
本文链接：https://www.haomeiwen.com/subject/bqunhrtx.html
网友评论