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 элементы будут отсортированы.

Параметры:

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

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

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