pytorch笔记：topk

1 介绍 

• 返回给定输入张量沿指定维度的最大的 k 个元素

• 如果没有指定 dim 参数，则默认选择输入张量的最后一个维度
• 如果 largest 设置为 False，则返回最小的 k 个元素。
• 返回一个命名元组 (values, indices)，其中包含每一行在指定维度 dim 上最大的 k 个元素的值和索引。
• 如果布尔选项 sorted 设置为 True，则确保返回的 k 个元素是有序的。

2 参数

3 举例

import torch
 
x=torch.arange(12).reshape(2,6)
 
x
'''
tensor([[ 0,  1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10, 11]])
'''
 
 
torch.topk(x,3)
'''
torch.return_types.topk(
values=tensor([[ 5,  4,  3],
        [11, 10,  9]]),
indices=tensor([[5, 4, 3],
        [5, 4, 3]]))
'''
 
 
 
torch.topk(x,3).indices
'''
tensor([[5, 4, 3],
        [5, 4, 3]])
'''