Use a pre-trained model to classify images in 50 categories of my choice

Question

As I said in the title, I would like to use a pre-trained ML model like InceptionV3 or similar to classify images from Instagram profiles.
I'm currently using PyTorch and I can correctly made inference on the models that I'm trying, but I'm not finding a nice way to "restrict" the outputs of the model to match my desired outcomes, that are only 50 categories of my choice.
I would like a lot to have some suggestions on how to do it.
Thank you in advance guys.

Answer 1

You can achieve the same by modifying two layers from the pretrained version as:

from torchvision.models import inception_v3

# load the pre-trained model
model = inception_v3(pretrained=True)

# Freeze the weights of the earlier layers since fine-tuning
for param in model.parameters():
    param.requires_grad = False

# Replace the last fully connected layer
inp_feat_lin = model.fc.in_features
model.fc = torch.nn.Linear(inp_feat_lin, 50)

# Replace the AuxLogits layer
inp_feat_aux = model.AuxLogits.fc.in_features
model.AuxLogits.fc = torch.nn.Linear(inp_feat_aux, 50)

# Set the `requires_grad` attribute to `True` for the last layer(s) since fine-tuning
model.fc.requires_grad = True
model.AuxLogits.fc.requires_grad = True

Add this to your class Definition of the model. You can read more about it in depth on the pytorch tutorials page itself.