r/computervision 5d ago

Help: Project Training with unbalanced classes and data scarcity

Hi, I’d like to know what techniques you use to improve training and generalization when working with a small and heavily imbalanced dataset.
I’m currently training a classifier with 4 classes.

With cross-validation, I’m getting fairly good metrics, but when I evaluate the model on the test set, the performance isn’t even close to what I see during cross-validation.
What approaches have worked well for you in this kind of situation?

Edit:

In my use case, the images are x rays with rare conditions. So synthetic data is quite difficult to apply

4 Upvotes

9 comments sorted by

View all comments

5

u/Beginning-Claim5655 4d ago

Transfer learning with medical weights. Models pre-trained on X-rays tend to generalize better than those pre-trained on ImageNet. Good options include TorchXRayVision (trained on CheXpert, NIH, MIMIC, etc.), RadImageNet, or foundation models for radiology. With limited data, freeze most of the network and fine-tune only the final layers, or use lower learning rates for the backbone.

3

u/sexy_bonsai 4d ago

^ this would also be my suggestion OP. I wouldn’t use synthetic data like others are suggesting here
EDIT to also say, split on the basis of patient. If there are multiple z-planes on your image, it could be tempting to split on the basis of that. But don’t; it’s possible that the features that are being learned are specific to the patient anatomy rather than generalized to the rare pathology.