Функция nn.top_k

Функция nn.top_k возвращает k наибольших элементов тензора вдоль последней оси. Первым параметром функция принимает входной тензор. Вторым параметром передается количество элементов k. Третьим параметром можно указать логическое значение sorted, определяющее, будут ли отсортированы результаты по убыванию.

Синтаксис

tf.nn.top_k(input, k, [sorted])

Пример

Давайте найдем 3 наибольших элемента в тензоре из списка чисел 1, 2, 3, 4, 5:

import tensorflow as tf

t = tf.constant([1, 2, 3, 4, 5])
res = tf.nn.top_k(t, k=3)
print(res)

Результат выполнения кода:

TopKV2(values=tf.Tensor([5 4 3], shape=(3,), dtype=int32), indices=tf.Tensor([4 3 2], shape=(3,), dtype=int32))

Пример

Давайте извлечем значения и индексы отдельно с помощью атрибутов values и indices:

import tensorflow as tf

t = tf.constant([10, 5, 20, 15, 30])
res = tf.nn.top_k(t, k=2)
print(res.values)
print(res.indices)

Результат выполнения кода:

tf.Tensor([30 20], shape=(2,), dtype=int32)
tf.Tensor([4 2], shape=(2,), dtype=int32)

Пример

Давайте применим функцию к двумерному тензору и найдем 2 наибольших элемента в каждой строке:

import tensorflow as tf

t = tf.constant([[1, 2, 3], [4, 5, 6]])
res = tf.nn.top_k(t, k=2)
print(res.values)
print(res.indices)

Результат выполнения кода:

tf.Tensor(
[[3 2]
 [6 5]], shape=(2, 2), dtype=int32)
tf.Tensor(
[[2 1]
 [2 1]], shape=(2, 2), dtype=int32)

Смотрите также

функцию in_top_k,
которая проверяет, входят ли элементы в k наибольших
функцию softmax,
которая применяет softmax к тензору
функцию log_softmax,
которая вычисляет логарифмический softmax
функцию l2_normalize,
которая нормализует тензор по L2-норме