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 -

Dataset -

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)


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.

Discussion (0)