Функция argmax

Функция argmax возвращает индекс наибольшего значения в тензоре вдоль заданной оси. Первым параметром функция принимает входной тензор. Вторым параметром передается ось axis, по которой нужно найти максимум. Третьим параметром можно указать тип выходных данных output_type.

Синтаксис

tf.argmax(input, axis=None, output_type=tf.int64)

Пример

Давайте найдем индекс максимального значения в одномерном тензоре:

import tensorflow as tf

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

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

tf.Tensor([4], shape=(1,), dtype=int64)

Пример

Давайте найдем индексы максимальных значений по столбцам для двумерного тензора:

import tensorflow as tf

t = tf.constant([[1, 2, 3], [4, 5, 6]])
res = tf.argmax(t, axis=0)
print(res)

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

tf.Tensor([1 1 1], shape=(3,), dtype=int64)

Пример

Давайте найдем индексы максимальных значений по строкам:

import tensorflow as tf

t = tf.constant([[1, 2, 3], [4, 5, 6]])
res = tf.argmax(t, axis=1)
print(res)

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

tf.Tensor([2 2], shape=(2,), dtype=int64)

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

функцию argmin,
которая возвращает индексы минимальных значений
функцию reduce_max,
которая вычисляет максимальное значение
функцию maximum,
которая возвращает поэлементный максимум
функцию top_k,
которая возвращает k наибольших значений и их индексы