torch.topk¶
- torch.topk(input, k, dim=None, largest=True, sorted=True, *, out=None)¶
返回给定
input
张量沿给定维度的前k
个最大元素。如果未提供
dim
,则选择 input 的最后一个维度。如果
largest
为False
,则返回前 k 个最小元素。返回一个名为元组的 (values, indices),其中包含 input 张量在给定维度 dim 上每行的前 k 个最大元素的 values 和 indices。
如果布尔选项
sorted
为True
,则确保返回的 k 个元素本身已排序- 参数
- 关键字参数
out (tuple, 可选) – (Tensor, LongTensor) 的输出元组,可以可选地提供用作输出缓冲区
示例
>>> 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]))