Función PyTorch torch.argmax

Pytorch torch 参考手册 Manual de referencia de Pytorch torch

torch.argmaxes una función en PyTorch que se utiliza para devolver el índice del valor máximo a lo largo de una dimensión.

Definición de la función

torch.argmax(input, dim, keepdim)

Ejemplos de uso

Ejemplo

import torch

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

# Índice del valor máximo global
print("Índice máximo global:", torch.argmax(x))

# Índice del valor máximo a lo largo de dim=0
print("Índice máximo a lo largo de dim=0:", torch.argmax(x, dim=0))

# Índice del valor máximo a lo largo de dim=1
print("Índice máximo a lo largo de dim=1:", torch.argmax(x, dim=1))

El resultado de salida es:

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

Pytorch torch 参考手册 Manual de referencia de Pytorch torch

Otras extensiones