Python PyTorch topk用法及代码示例

本文简要介绍python语言中 torch.topk 的用法。

用法:

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

参数：

• input(Tensor) -输入张量。

• k(int) -“top-k”中的k

• dim(int,可选的) -要排序的维度

• largest(bool,可选的) -控制是否返回最大或最小元素

• sorted(bool,可选的) -控制是否按排序顺序返回元素

关键字参数：

out(tuple,可选的) -(Tensor, LongTensor) 的输出元组，可以选择用作输出缓冲区

返回给定 input 张量沿给定维度的 k 最大元素。

如果没有给出dim，则选择input的最后一个维度。

如果largest 是False，则返回k 最小元素。

返回 (values, indices) 的命名元组，其中 indices 是原始 input 张量中元素的索引。

布尔选项 sorted if True 将确保返回的 k 元素本身是排序的

例子：

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