Custom DataLoader for two datasets with matching labels

1

I am trying to write a custom DataLoader to combine two image datasets. These datasets have different lengths but the same label classes. Dataset A is relatively small, while Dataset B is relatively big. The procedure would look like this:

  1. Sample a batch from Dataset A with label x1.
  2. Sample a batch from Dataset B with label x1.
  3. Repeat the above steps for all remaining labels x2, x3…
  4. Sample another batch from Dataset A with label x1.
  5. Sample another batch from Dataset B with label x1.
  6. And so on.
  7. Once all the images from Dataset A are exhausted, I would like to start over for Dataset A.
  8. Once all the images from Dataset B are exhausted, the procedure stops.

I am aware of ConcatDataset, but it doesn’t really suit my purpose. One idea would be to use Subset to construct different iterators for different classes. This is obviously not practical. Is there another way to go about this I am missing?