What is the predicted output label from a PyTorch model?

Question

I am currently training a ResNet18 model with a custom optimizer in PyTorch.

I am using CrossEntropyLoss() and the ResNet18 model from PyTorch. In tensorflow the outputs are of the desired shape, but in pytorch it is necessary to find the argmax of the predicted labels in order to find the accuracy.

If my batch size = 64 with the resnet model, why is the predicted label of shape [64, 1000]?

What do the 1000 values correspond to?

ayandas · Accepted Answer · 2021-10-13 17:42:41Z

1

The predicted quantity is not "label", it is the probability (soft score) of the input being one of 1000 classes.

The output of (64, 1000) contains a 1000 length vector for each input in a batch. If you want discrete labels (i.e. 0 to 999), perform an argmax over it

labels = torch.argmax(output, 1)

By argmax over each probability vector, we compute which class (among 1000) has the highest probability for the input.

answered Oct 13, 2021 at 17:42

ayandas

2,3181 gold badge15 silver badges29 bronze badges

Sign up to request clarification or add additional context in comments.

5 Comments

user16573587 Over a year ago

I am using cifar10 which doesn't have 1000 classes.

ayandas Over a year ago

Then you have to make the last layer 10 and train the network

user16573587 Over a year ago

does the torchvision.models.resnet18() model not work for this? Do I need to make it myself?

user16573587 Over a year ago

I just switched to a different model and it is working. Thank you

ayandas Over a year ago

yes it does, just pass num_classes=10 in the constrcutor, i.e. model = ResNet(num_classes=10)

Collectives™ on Stack Overflow

What is the predicted output label from a PyTorch model?

1 Answer 1

5 Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

5 Comments

Related