Retrain a pretrained model after adding Softmax layer

vision
1

Hi all, I wanted to know if it is needed to retrain the pretrained pytorch model if I wish to add a softmax layer at the end to get the class probabilities in case of classification problem?
Eg.

net = torchvision.models.resnet50(pretrained = True)
net.fc = nn.Sequential(nn.Linear(in_features=2048, out_features=1000))
net.fc = nn.Sequential(*list(net.fc)+[nn.Softmax(1)])

I am exporting this modified resnet50 pretrained model as onnx. On deployment on my browser, the class probabilities are being returned as 0% for top-4 classes. Previously using the default pretrained model it was coming as 1600% and so as there isn’t a context of class probabilities in the output layer in default pretrained resnet50.
Is there a way I can introduce softmax without needing to train the network again?
Thanks!!

2

You could simply add a softmax layer to the model if you need to see the probabilities.
The output and the predicted class won’t change, as the max logit is also the max probability.
To get the predicted class you could use torch.aargmax(output, 1) with output being the logits or probabilities.

1 Like