This python demo notebook shows how to fine-tune a 1000 class output network to a 2 class output network. The dataset I have consists of 60 data points, 30 for each class (dog and cat images), and 24 data points for validation dataset, 12 for each class.
Python Notebook -
https://github.com/rishiraj824/MachineLearningAlgorithms/blob/master/resnet/resnet34-finetuning.ipynb
Dataset -
https://github.com/rishiraj824/MachineLearningAlgorithms/tree/master/resnet/data
Basic requirements for training a large network like resnet is that you have to choose a valid loss function and an optimizer. For this example, I have chosen a Cross-Entropy Loss which is minimized by my model. This is useful as it already comes with a Softmax layer in it, in the pytorch library.
I am using the SGD (Stochastic Gradient Algorithm) to solve for the minimization and to optimize the model. And then at the end of the trained model, I add a Linear layer which changes the feature sets to 2 classes
num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 2)
Observations:
Plotting the Training Loss and the Validation Loss with the iteration shows us how quickly it converged. This was bound to happen as this is a large network and the dataset we used was very small. But since our main objective was to fine-tune the network we were able to achieve that. This code would work for any number of data size. I hope this helped to get a brief idea of how fine-tuning works.
Top comments (0)