torch.argmax

torch.argmax(input) → LongTensor

Возвращает индексы максимального значения всех элементов в input тензоре.

Это второе значение, возвращаемое функцией torch.max(). Смотрите её документацию для точного определения семантики этого метода.

Примечание

Если существует несколько максимальных значений, возвращаются индексы первого максимального значения.

Параметры

input (Тензор) – входной тензор.

Пример:

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 1.3398,  0.2663, -0.2686,  0.2450],
        [-0.7401, -0.8805, -0.3402, -1.1936],
        [ 0.4907, -1.3948, -1.0691, -0.3132],
        [-1.6092,  0.5419, -0.2993,  0.3195]])
>>> torch.argmax(a)
tensor(0)

torch.argmax(input, dim, keepdim=False) → LongTensor

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

Это второе значение, возвращаемое функцией torch.max(). Смотрите её документацию для точного определения семантики этого метода.

Параметры

• input (Тензор) – входной тензор.
• dim (int) – размерность для сокращения. Если None, возвращается argmax сплющенного входного тензора.
• keepdim (bool) – сохраняется ли размерность выходного тензора. Игнорируется, если dim=None.

Пример:

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 1.3398,  0.2663, -0.2686,  0.2450],
        [-0.7401, -0.8805, -0.3402, -1.1936],
        [ 0.4907, -1.3948, -1.0691, -0.3132],
        [-1.6092,  0.5419, -0.2993,  0.3195]])
>>> torch.argmax(a, dim=1)
tensor([ 0,  2,  0,  1])