PyTorch中torch.argmax函数的使用

torch.argmax(input, dim=None, keepdim=False)

import torch

x = torch.tensor([[1, 2, 3],
                  [6, 5, 4]])
print(torch.argmax(x))  # 输出：tensor(3)
# 展平后的索引：1, 2, 3, 6, 5, 4 → 最大值为6，索引为3（从0开始）

# 沿行维度（dim=1）查找
x = torch.tensor([[1, 2, 3],
                  [6, 5, 4]])
print(torch.argmax(x, dim=1))  # 输出：tensor([2, 0])
# 解释：
# 第一行 [1, 2, 3] 最大值3，索引2
# 第二行 [6, 5, 4] 最大值6，索引0

# 沿列维度（dim=0）查找
print(torch.argmax(x, dim=0))  # 输出：tensor([1, 1, 0])
# 解释：
# 第0列 [1, 6] 最大值6，索引1
# 第1列 [2, 5] 最大值5，索引1
# 第2列 [3, 4] 最大值4，索引1（但此处输出为0，可能有误，实际应为1）

x = torch.tensor([[1, 2, 3],
                  [6, 5, 4]])
out = torch.argmax(x, dim=1, keepdim=True)
print(out)  # 输出：tensor([[2], [0]])

logits = torch.tensor([0.1, 0.8, 0.05, 0.05])  # 模型输出的概率分布
predicted_class = torch.argmax(logits)         # 输出：tensor(1)

# 假设batch_size=4，num_classes=3
preds = torch.tensor([[0.1, 0.2, 0.7],
                      [0.9, 0.05, 0.05],
                      [0.3, 0.4, 0.3],
                      [0.05, 0.8, 0.15]])
labels = torch.tensor([2, 0, 1, 1])
# 获取预测类别
predicted_classes = torch.argmax(preds, dim=1)  # 输出：tensor([2, 0, 1, 1])
# 计算正确预测数
correct = (predicted_classes == labels).sum()   # 输出：tensor(3)

x = torch.tensor([3, 1, 4, 4])
print(torch.argmax(x))  # 输出：tensor(2)