如何将输出张量转为最大值索引张量并计算其与目标张量的准确率？

1. 转换输出张量为最大元素索引张量 #

直接使用PyTorch的torch.argmax()函数，指定dim=1即可对每一行（每个样本）取最大元素的索引，得到你需要的一维张量：

import torch

# 你的输出张量
output = torch.tensor([[0.1834, 0.8166],
                       [0.3031, 0.6969],
                       [0.3104, 0.6896],
                       [0.3065, 0.6935],
                       [0.3060, 0.6940],
                       [0.2963, 0.7037],
                       [0.2340, 0.7660],
                       [0.2302, 0.7698],
                       [0.2581, 0.7419],
                       [0.2081, 0.7919]], requires_grad=True)

# 转换为最大元素索引张量，dim=1表示按行取最大值索引
pred_indices = output.argmax(dim=1).float()
# 输出结果：tensor([1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])

2. 计算与目标张量的准确率 #

假设你有对应的目标张量target（维度与pred_indices一致），通过以下步骤计算准确率：

• 比较预测索引与目标值是否相等，得到布尔张量
• 对布尔张量求均值，即为准确率（True会被视为1，False视为0）

示例代码：

# 假设目标张量（示例数据，根据实际情况替换）
target = torch.tensor([1., 1., 0., 1., 1., 1., 0., 1., 1., 1.])

# 计算准确率
accuracy = (pred_indices == target).float().mean()
print(f"准确率: {accuracy.item():.4f}")
# 示例输出：准确率: 0.8000

说明 #

• 如果目标张量是整数类型，无需转换pred_indices为float，直接比较即可
• 若需要保留梯度，注意argmax会中断梯度传播，但准确率计算一般在验证阶段，不需要梯度，所以可以正常使用

内容的提问来源于stack exchange，提问作者shuai wang