Функция PyTorch torch.argmax

Pytorch torch 参考手册 Справочное руководство PyTorch torch

torch.argmaxЭто функция в PyTorch, используемая для возврата индекса максимального значения вдоль размерности.

Определение функции

torch.argmax(input, dim, keepdim)

Пример использования

Пример

import torch

x = torch.tensor([[1, 3, 2], [4, 1, 3]])

# Индекс глобального максимального значения
print("Индекс глобального максимума:", torch.argmax(x))

# Индекс максимального значения вдоль dim=0
print("индекс максимума для dim=0:", torch.argmax(x, dim=0))

# Индекс максимального значения вдоль dim=1
print("индекс максимума для dim=1:", torch.argmax(x, dim=1))

Результат вывода:

全局最大索引: tensor(4)
dim=0 最大索引: tensor([1, 0, 1])
dim=1 最大索引: tensor([1, 0])

Pytorch torch 参考手册 Справочное руководство PyTorch torch

Другие расширения