如何得到预测概率?

8 pytorch

此代码从模型中获取10值。
如果我想得到预测的概率
我应该改变哪一行?

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()
Run Code Online (Sandbox Code Playgroud)

Sha*_*hai 14

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

import torch.nn.functional as nnf

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

top_p, top_class = prob.topk(1, dim = 1)
Run Code Online (Sandbox Code Playgroud)

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