CODEWITHSUNDEEP
deep-learningprobabilitypytorch
龚世泽
龚世泽

Reputation: 21

How to get probabilities when using Pytorch's densenet?

I want to do a binary classification and I used the DenseNet from Pytorch.

Here is my predict code:

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

And here is the output:

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

I want to get the probabilities of each class. What should I do?

I have noticed that torch.nn.Softmax() could be used when there are many categories, as discussed here.

Upvotes: 1

Views: 693

Answers (1)

Ben D
Ben D

Reputation: 21

import torch.nn as nn

Add a softmax layer to the classifier layer: i.e. typical:

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(dim=1))

Upvotes: 1

Related Questions