如何获得预测概率？

Question

此代码从模型获取1或0值。
如果我想获得预测的可能性
我应该更改哪一行？

from torch.autograd import Variable
results = []
#names = []
with torch.no_grad():
    model.eval()
    print('===============================================start')
    for num, data in enumerate(test_loader):
        #print(num)
        print("=====================================================")

        imgs, label = data
        imgs,labels = imgs.to(device), label.to(device)
        test = Variable(imgs)
        output = model(test)
        #print(output)
        ps = torch.exp(output)
        print(ps)
        top_p, top_class = ps.topk(1, dim = 1)
        results += top_class.cpu().numpy().tolist()



model = models.resnet50(pretrained=True)
model.fc = nn.Linear(2048, num_classes)
model.cuda()

Answer 1

模型通常输出原始预测logit。要将它们转换为概率，您应该使用softmax函数

import torch.nn.functional as nnf

# ...
prob = nnf.softmax(output, dim=1)

top_p, top_class = prob.topk(1, dim = 1)

新变量top_p应该为您提供前k个类别的概率。