详解torch.topk

torch.tok()

 

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

参数：

         

• input (Tensor) – 输入张量（必须是torch的张量类型）

• k (int) – 需要返回的topk的个数

• dim (int, optional) – the dimension to sort along，默认为最后一维

• largest (bool, optional) –是否返回最大的topk个元素，默认为True

• sorted (bool, optional) –是否对返回元素进行排序，默认为True，即降序排列

• return：返回值为k个在 一定维度上最大或最小的元素和它们的索引

如果dim参数未指定，则默认最后一维

如果largest=False，则返回最小的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]))

inference：

1.https://pytorch.org/docs/stable/generated/torch.topk.html?highlight=torch%20topk#torch.topk