How to train on a multiple object dataset for 2D case (1-shot segmentation for multi-class dataset)

[jacoblam3112]

Hi Thank you for sharing your code.

How can I train on multiple datasets ( or a single dataset with multiple class of images) simultaneously while making sure different classes are not mixed in the memory bank (image flow) during training ?

Should I do something like following for example:
Put 5 images of dataset A(1 prompt image and 4 training images) calculate the gradients, and then empty the memory bank
Put 5 images of dataset B calculate gradients (do gradient accumulation), empty the memory bank
do optimiser and then repeat
?