pytorch笔记：topk

1 介绍 

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

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

2 参数

input (Tensor)	输入张量
k (int)	表示“top-k”的 k 值
dim (int, 可选)	要排序的维度
largest (bool, 可选)	控制是返回最大还是最小的元素
sorted (bool, 可选)	控制是否以排序的顺序返回元素。

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]])
'''