torch.argmax()

torch.argmax 是 PyTorch 中的一个函数，它用于返回给定张量中最大值的索引。这个函数通常用于多分类问题中，用来找出模型预测最可能的类别。

简单来说就是返回给定张量中最大值的索引。

实例：

import torch

# 假设这是模型对一个样本的预测输出，形状为 [1, 10]
# 每一列代表模型预测样本属于对应花的概率
predict = torch.tensor([[0.1, 0.05, 0.2, 0.3, 0.1, 0.05, 0.05, 0.05, 0.05, 0.0]])

# 使用 argmax 找到概率最高的花的索引
predict_cla = torch.argmax(predict)
print(predict_cla)

# 将结果转换为 NumPy 数组
predict_cla_np = predict_cla.numpy()
print(predict_cla_np)

# 花的名称列表
flower_names = ["玫瑰", "牡丹", "菊花", "月季", "百合", "郁金香", "康乃馨", "向日葵", "紫罗兰", "风信子"]

# 根据索引找到对应的花的名称
predicted_flower = flower_names[predict_cla_np]

print("预测的花的名称:", predicted_flower)

输出：

tensor(3)
3
预测的花的名称: 月季