Efficient Selective Pre-Training for Imbalanced Fine-Tuning Data in Transfer Learning
Shaif Chowdhury, Sadia Nasrin Tisha, Mushfika Rahman, Greg Hamerly · 2024
Neural networks are often pre-trained on a large source dataset and then fine-tuned on a smaller target dataset. Although pre-training on large-scale datasets is very useful, it has a few disadvantages, such as (1) high training cost and (2) domain mismatch where pre-training on a less-related source might lead to poor results in a target model. Examples of this are areas like underwater imaging, medical imaging, microscopic imaging, etc. Many datasets in these domains also have class imbalance which makes transfer learning less effective. In this paper, we propose an efficient method for selective pre-training, i.e. selecting relevant subsets from a pre-training dataset. Fine-tuning with our method gives better accuracy while increasing training efficiency. We validate our technique with selective pre-training on ImageNet21k and ImageNet1k datasets, and fine-tuning on tasks like image classification and image segmentation. We conduct experiments on several imbalanced datasets and compare our performance with full pre-training as well as other state-of-the-art methods to handle class imbalance. On imbalanced CIFAR-10 we get an accuracy of 77% with pre-training on 500k images of ImageNet1k compared to 74% for full pre-training on ImageNet.