How to modify the final FC layer based on the torch.model

34

I seem to be a bit late to the party, but I just wanted to share a supplementary solution for this topic which should work for any VGG:

model = torchvision.models.vgg19_bn(pretrained=True)
n_feats = 1 #   whatever is your number of output features
last_item_index = len(model.classifier)-1
old_fc = model.classifier.__getitem__(last_item_index )
new_fc = nn.Linear(in_features=old_fc.in_features, out_features= n_feats, bias=True)
model.classifier.__setitem__(last_item_index , new_fc)
5 Likes