torch.topk

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

Возвращает k наибольших элементов заданного тензора input вдоль указанного измерения.

Если dim не задано, выбирается последнее измерение input.

Если largest равно False, возвращаются k наименьших элементов.

Возвращается именованный кортеж из (values, indices) с values и indices наибольших k элементов каждой строки тензора input в заданном измерении dim.

Булев параметр sorted, если он равен True, гарантирует, что возвращаемые k элементы будут отсортированы

Примечание

При использовании torch.topk порядок индексов элементов с одинаковыми значениями не гарантируется и может различаться при разных вызовах.

Параметры:

• input (Tensor) – входной тензор.
• k (int) – значение k в операции «top-k»
• dim (int, optional) – измерение, вдоль которого выполняется сортировка
• largest (bool, optional) – определяет, возвращать наибольшие или наименьшие элементы
• sorted (bool, optional) – определяет, возвращать ли элементы в отсортированном порядке

Именованные аргументы:

out (tuple, optional) – выходной кортеж (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]))