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 (Тензор) – входной тензор.
• k (целое число) – значение k в «топ-k».
• dim (целое число, необязательно) – измерение для сортировки.
• largest (логическое значение, необязательно) – определяет, возвращать ли наибольшие или наименьшие элементы.
• sorted (логическое значение, необязательно) – определяет, возвращать ли отсортированные элементы.

Ключевые аргументы

out (кортеж, необязательно) – выходной кортеж (Тензор, 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]))