使用Pytorch的密网时如何获得概率？

Question

我想进行二进制分类，并且使用了Pytorch的DenseNet。

这是我的预测代码：

densenet = torch.load(model_path)
densenet.eval()
output = densenet(input)
print(output)

这是输出：

Variable containing:
54.4869 -54.3721
[torch.cuda.FloatTensor of size 1x2 (GPU 0)]

我想获得每个班级的概率。我该怎么办？

我注意到torch.nn.Softmax()可以在许多类别中使用，如here所述。

Answer 1

[将softmax图层添加到分类器图层：即典型：

num_ftrs = model_ft.classifier.in_features
model_ft.classifier = nn.Linear(num_ftrs, num_classes)

 updated:

model_ft.classifier = nn.Sequential(nn.Linear(num_ftrs, num_classes), 
nn.Softmax(num_classes, dim=1))