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

 

你可能感兴趣的:(AI)