How to switch dataloader every n training steps
Hi everyone,
In my current setup, I would like to change the dataloader during a training epoch:
This is what I would like to achieve:
step 1.Train on dataset 1 for n batches
step 2.Train on dataset 2 for n batches
step 3.Go to step 1
I found this solution on the old forum but this only switches the dataset after each epoch.
Here is my current attempt at switching it every n batches:
class SimpleModule(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = ...
        self.batch_size = ...
        self.change_every_n_batch = 20
    
    def train_dataloader(self):
        self.current_dataset = (self.global_step // self.change_every_n_batch) % 2
        if self.current_dataset == 0:
            dataset = Dataset1()
        elif self.current_dataset == 1:
            dataset = Dataset2()

        dataloader = DataLoader(dataset, batch_size=self.batch_size)
        return dataloader
    
    def on_train_batch_end(self, outputs, batch, batch_idx):
        new_dataset = (self.global_step // self.change_every_n_batch) % 2
        if new_dataset != self.current_dataset:
            self.trainer.reset_train_dataloader(self)
train_dataloader() is called as expected every 20 batches by on_train_batch_end() but the returned dataloader does not seem to be used during the training loop.
Any idea what could be going wrong? Or do you have a solution for what I want to achieve?
Thanks!
hey @matprst !
you can set:

limit_train_batches=n. This will ensure that every training epoch will progress for only n batches
reload_dataloaders_every_n_epochs=1. this will ensure that train dataloader is reloaded after every epoch.

and inside train_dataloader, flip the dataloader on each reload. something like:
def train_dataloader(self):
    if self.some_flag:
        dataset = Dataset1()
    else:
        dataset = Dataset2()

    self.some_flag = not self.some_flag

    return DataLoader(dataset, batch_size=self.batch_size)