如何从已训练的PyTorch模型中获取0-1或0-100区间的置信度得分

你当前拿到的输出是模型全连接层直接输出的原始logits，数值范围无界，只需要对logits做Softmax归一化即可得到0-1区间的置信度，乘以100即可得到0-100区间的得分。

代码修改步骤 #

1. 首先导入PyTorch的函数库：

import torch.nn.functional as F

2. 修改预测函数，在取最大值前先对输出做Softmax处理：

def predict_allCharacters(imgs):
    output = model(imgs)
    # 沿类别维度做Softmax，将所有类别的输出归一化到0-1区间，总和为1
    probs = F.softmax(output, dim=1)
    conf, classes = torch.max(probs, 1)
    # 如果需要0-100区间的置信度，取消下行注释即可
    # conf = conf * 100
    class_names = '0123456789'
    return conf, class_names[classes.item()]

注意事项 #

• 如果你训练模型时用的是CrossEntropyLoss，损失函数内部已经集成了Softmax计算，推理阶段单独加Softmax不会影响模型精度，无需重新训练。
• 你当前模型的全连接层输出维度为37，但class_names仅定义了10个数字类别，如果你实际训练的类别包含字母等其他分类，记得同步修改class_names的内容，避免类别映射错误。

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