Torchvision models dont have softmax layer？

vision
1

In the torchvision models vgg code：https://github.com/pytorch/vision/blob/master/torchvision/models/vgg.py

there is not a softmax layer in this code ,but in the original paper the last layer is a softmax layer，why dose it happen？
if I use pretrained vgg model to recognize a picture，do I need to add a softmax layer in the last of the model to get each class’s probability？

5 Likes
2

torchvision models were trained usingnn.CrossEntropyLoss which consists of nn.LogSoftmax and then nn.NLLLoss.
that’s why there is no softmax layer at the end.
So you only need to use nn.CrossEntropyLoss() in your model

9 Likes
3

thanks for your reply！
when in the test mode，do I need to add a nn.Softmax layer after the model to get the probabilities of each class?

4

if you want the probability for each class then yes.
if you want only the predicted class then you don’t need it you can only use torch.max()

5 Likes