Hello Marsggbo!



 marsggbo:

I have a imbalanced dataset, for some classes, there are 3000 images, but for some classes, there are only 40 images for each class. So my idea is to train a model A on classes which have more images, while train a model B on classes which have less images.
Suppose two models perform well on each data, now how can I merge  these two models to predict all images?


I don’t think that there is a particularly good way to merge
the models as you suggest.
The intuition is as follows:
Let’s assume – as is likely the case – that the early layers
of your trained  A and B models are both extracting essentially
the same low-level features that are appropriate for use by
the later layers of the two models for performing the A and
B classifications.  Even so, because the two models were trained
independently, the particular way in which those features are
encoded in the early layers will happen to be different, and
the A-model low-level feature encoding won’t happen to match
what the later layers of the B-model are expecting.
I think it makes more sense to use any of various established
techniquies for dealing with unbalanced data.
Your supposition that your B-model performs well on your B data
implies that you have enough B data, even if it’s a lot less
than your A data.  If you simply don’t have enough B data, things
get harder.
One approach would be to sample your B data more heavily so that
you train your model with balanced samples.
To be concrete, let’s say that you have 5 A classes, each with
3000 sample images, and 5 B classes, each with only 40 images,
for 10 classes total.  If you were to use a batch size of 50
for each training step, you could force your (otherwise randomly
sampled) batches to have exactly 5 images from each class.  Of
course, as you train, you’ll start seeing repeats of your B-class
images much sooner than A-class repeats.  (But that’s okay.)
You don’t have to force your batches to have exactly 5 images
from each class.  You could sample so that you “probably” have
5 images.  You might take a look at
torch.utils.data.sampler.WeightedRandomSampler
for doing this kind of thing.
It could also be better to not fully reweight your B data.  For
example, you could partially reweight your B data so that a
typical batch has, say, 40 A images and 10 B images.  (What
scheme works best depends on the details of your problem.)
Best regards.
K. Frank