Multiple dataloaders in training_step() and use them separately

[thangld201]

Hi, I’m figuring out how to use multiple dataloaders in training_step() of LightningModule. Currently, I pass a list of dataloaders in trainer.fit(), it will return the list of batches, each from a dataloader simultaneously. However, my use case differs in that I would want to process each batch from each dataset sequentially.

For example, I have three datasets. For step i, I receive a batch from dataset 0, update my model. For step i+1, I receive a batch from dataset 1 and update my model. For step i+2, I get a batch from dataset 2 and update my model. The process repeats until all samples are iterated.

How can I implement this in Pytorch Lightning ? Are there already supports for this ? I would be happy to dive in myself, but I don’t know where to start.

[changspencer]

Hi @thangld201, this is an interesting use case.

Before reading the proposition below, you may want to answer a question: why do you need to have it iterate each dataset in a separate batch? Why not do what you're suggesting and feed dataset 1, compute loss, BP, feed dataset 2, compute loss, BP,..., feed dataset num_datasets, compute loss, and BP? Is it for logging purposes?

If it's more nuanced that the above, read on. 🙂

---

Ultimately, you may not want to use the following, but there could be one or two good ideas.

A brute force approach may be to load all the datasets into a single Dataset object that would return the particular dataset based on the index that is input to the __getitem__ method for these objects. Reminder: the __getitem__ method is given as __getitem__(self, index).

For example, suppose that you have num_datasets to pull from. Inside __getitem__, put some conditional that looks at the value of data_idx = index % num_datasets and chooses the dataset corresponding to that value. It could look something like the following:

if data_idx == 0:
    return dataset1_preproc(index)
elif data_idx == 1:
    return dataset2_preproc()
elif data_idx == 2:
    return dataset3_preproc()
else:
    # Whatever your default case would be if there's something wrong with the indexing.

The downside to this approach is that - because you may look to "repurpose" the index - multiple problems come up:

• How do you make the dataloader properly view the number of possible batches?
• How do you do shuffled sampling during training (may have to implement this manually)?
  • A potential way around the indexing problem is to hold an "internal index" that iterates over each datasets every time the dataset object is called.

[AndrewAnnex]

@changspencer I am having a similar issue so want to post here and not open a new discussion

After following the info about
arbitrary iterable support (https://lightning.ai/docs/pytorch/stable/data/iterables.html#multiple-iterables) in a small DataModule I wrote, in my training_step call the batch is now a list of batch dicts, and from the docs it's unclear what the user should do in that situation given the expected return type. It seems that just doing that alone is probably passing in alternating batches given the dataloader_idx, but I feel like I shouldn't have to handle the batch as a list and that somewhere outside of training_step Lightning should just be iterating through the batches for me allowing me to leave my code unchanged.

I don't want to fit my model sequentially as I feel that'd just leave to the model forgetting/overfitting with class imbalance I believe is present between the two datasets, but that's really just a hunch so if that is really the correct way to do it let me know.

[tungts1101]

@changspencer To answer why I want to handle each data loader separately, the reason is that in the robustness test, I want a result for each corruption to be logged separately. Can we do anything in the test_step() method?

[iliasslasri]

hey, do you have any update on that ? I'm interested!