pytorch 中的topk函数

pytorch中topk() 函数用法

 

1. 函数介绍

最近在代码中看到这两个语句

maxk = max(topk)
_, pred = output.topk(maxk, 1, True, True)

这个函数是用来求output中的最大值或最小值，返回两个参数：其一返回output中的最大值（或最小值），其二返回该值的索引。

 

2. topk()函数原型：

（具体的用法参考pytorch官方中文文档：http://www.mamicode.com/info-detail-2217311.html）

torch.topk(input, k, dim=None, largest=True, sorted=True, out=None) -> (Tensor, LongTensor)

参数：

• input (Tensor) – 输入张量
• k (int) – “top-k”中的k
• dim (int, optional) – 排序的维
• largest (bool, optional) – 布尔值，控制返回最大或最小值
• sorted (bool, optional) – 布尔值，控制返回值是否排序
• out (tuple, optional) – 可选输出张量 (Tensor, LongTensor) output buffers

例如：

k=1时，求的就是top-1的值

dim=1时，按行求最大或最小值

largest=True时，为求最大值，largest=False时，求最小值

 

3. 实例：

>>> import torch
>>> output=torch.randn((4, 2))
>>> output
tensor([[-0.3249,  1.0216],
        [-0.6855, -0.0272],
        [ 0.4624,  0.1392],
        [ 1.8406, -0.3436]])
>>> _,pred=output.topk(1, 1, True, True)     # 求top1
>>> _
tensor([[ 1.0216],
        [-0.0272],
        [ 0.4624],
        [ 1.8406]])
>>> pred
tensor([[1],
        [1],
        [0],
        [0]])