I have a custom dataset with 10 classes and I am using a pre-trained resnet18 model from torch-vision. I can clearly see it's over-fitting because: the model is trained for 75 epochs with a batch size of 4 and from epoch 30 the validation accuracy stops increasing and the training accuracy keeps increasing.

Things I did to improve generalisation performance and reduce overfitting:

  1. I normalized my data with a calculated mean and std on my training data.
  2. I added random rotation, random horizontal flip and random vertical flip as data augmentations.

The above were the results after all of these, I can't think of other regularization techniques that produce regularization effects without changing the architecture of the model which will completely contradict the purpose of a pre-trained network, and I am using one because I only have around 300 images in my each class.

  • $\begingroup$ My first question to this is that there are other pre-trained networks available for computer vision (which what I am assuming you are doing, given the model you referred to) which you could use instead? Usually during the ore-training phase, these models are trained on vast quantities of data, which span a variety of classes, so typically data is not the issue which causes poor generalisation performance. $\endgroup$ – shepan6 Jul 12 '20 at 7:08
  • 1
    $\begingroup$ Which layers did you re-train? $\endgroup$ – Sammy Jul 12 '20 at 8:11
  • $\begingroup$ @Sammy i trained the layer4 and the last fc layer $\endgroup$ – Aakash Kaushik Jul 12 '20 at 10:33
  • $\begingroup$ @shepan6 yes i know that there are other networks and i have tried AlexNet, vgg16 with batch_normalization, resnet18,34,50,101, Resnet wide with different configs and even the resnext. $\endgroup$ – Aakash Kaushik Jul 12 '20 at 10:34
  • $\begingroup$ Why Layer4. You should start with the last layer. Any specific reason? $\endgroup$ – 10xAI Jul 12 '20 at 11:28

In transfer learning there are two parameters which influence the basic setup to go for:

  1. Size of new dataset
  2. Similarity of new dataset to dataset of pre-trained model

When your dataset is small the problem is that high capacity pre-trained models can easily overfit if you re-train too many layers. And since you re-trained multiple layers this could be an issue here.

Instead, try the following two options:

  • Re-train only the last fully connected layer. This is the general approach if your data is similar to the data which the pre-trained model was trained on.
  • Remove the conv. layers towards the end of the pre-trained model and re-train only the new fully connected layer. For ResNet18 you could try tossing conv4 and 5, for example. This is the general approach in transfer learning for small datasets if your data is not very similar to the original data and you only want to use layers for lower level features since the datasets do not have similar higher level features.

If you are interested in a theoretical reading on the topic, the paper "How transferable are features in deep neural networks?" could be of interest.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.