DEV Community

Rishi Raj
Rishi Raj

Posted on

Fine Tuning Resnet34 for a 2 class, classification

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.

Heroku

This site is built on Heroku

Join the ranks of developers at Salesforce, Airbase, DEV, and more who deploy their mission critical applications on Heroku. Sign up today and launch your first app!

Get Started

Top comments (0)