Deep Learning with PyTorch

Fine-tuning VGG16

The VGG16 model is trained to classify 1,000 categories, but not trained to classify dogs and cats. So, we need to change the output features of the last layer to 2 from 1000. The following code snippet does it:

vgg.classifier[6].out_features = 2

The vgg.classifier gives access to all the layers inside the sequential model, and the sixth element will contain the last layer. When we train the VGG16 model, we only need the classifier parameters to be trained. So, we pass only the classifier.parameters to the optimizer as follows:

optimizer =   optim.SGD(vgg.classifier.parameters(),lr=0.0001,momentum=0.5)