torch.topk

函数定义

torch.topk(input, k, dim=None, largest=True, sorted=True, *, out=None)

对于给定的输入 张量input，沿着给定的维度，返回k个最大元素。

一个命名元组(values,indices)将会被返回，这里的indices是返回的元素在原始的input张量中的indices。

函数参数　

• input (Tensor) – the input tensor.

• k (int) – the k in “top-k”

• dim (int, optional) – the dimension to sort along，If dim is not given, the last dimension of the input is chosen.

• largest (bool, optional) – controls whether to return largest or smallest elements

• sorted (bool, optional) – controls whether to return the elements in sorted order

• out (tuple, optional) – the output tuple of (Tensor, LongTensor) that can be optionally given to be used as output buffers

例子

>>> x = torch.arange(1., 6.)
>>> x
tensor([ 1.,  2.,  3.,  4.,  5.])
>>> torch.topk(x, 3)
torch.return_types.topk(values=tensor([5., 4., 3.]), indices=tensor([4, 3, 2]))