How to switch data loaders between epochs while using multiple optimizer

[icedpanda]

Hi,

I have two datasets and each dataset will have a corresponding optimizer.

Right now I am updating the model parameters per batch per dataset. But how can I iterate datasetA for an epoch first and then iterate an epoch for datasetB?

• Train an epoch on the datasetA using optimizer A
• Train an epoch on the datasetB using optimizer B

This is what I have for iterating multiple datasets per batch.

class Mymodel(LightningModule):
  def configure_optimizers(self):
      optimizer = optim.Adam(
          filter(
              lambda p: p.requires_grad,
              self.parameters()),
          lr=self.lr,)
      scheduler = optim.lr_scheduler.ReduceLROnPlateau(
          optimizer,
      )
      scheduler_config = {
          "scheduler": scheduler,
          "monitor": "val/rec/epoch_loss",
          "interval": "epoch",
          "frequency": 1,
      }
      pre_optimizer = optim.SGD(
            filter(
                lambda p: p.requires_grad,
                self.parameters()),
            lr=self.lr_for_other_modules,)
      optimizers = [optimizer, pre_optimizer]
      return {"optimizer": optimizer, "lr_scheduler": scheduler_config}
  
  def training_step(self, batch, batch_idx, optimizer_idx):
       if optimizer_idx == 0:
          # do forward on datasetA
          inputs = batch["datasetA"]
        if self.is_use_some_modules and optimizer_idx == 1:
          # do something else 
          inputs = batch["datasetB"]
        return ...

class Datamodule(LightningDataModule):
  def train_dataloader(self):
      ...
      return {"datasetA": A_loader, "datasetB": B_loader}

I found a similar question here #3336 and the suggested solution is reloading data loaders every epoch.
I wonder if it is safe to do this as I have multiple optimizers.

class Datamodule(LightningDataModule):
  def train_dataloader(self):
      if self.current_epoch % 2 == 0:
          A_loader = ...
          return A_loader
      else:
          B_loader = ...
          return B_loader

class Mymodel(LightningModule):
  def training_step(self, batch, batch_idx, optimizer_idx):
      if self.current_epoch % 2 == 0 and optimizer_idx == 0:
           inputs = batch["datasetA"]
      elif self.current_epoch % 2 == 1 and optimizer_idx == 1:
          inputs = batch["datasetB"]
      return ...

[rohitgr7]

one simple solution is to use reload_dataloaders_every_epoch=1 and configure training_step based on current_epoch.

def training_step(self, batch, batch_idx, optimizer_idx):
    if (self.current_epoch % 2 == 0 and optimizer_idx == 0) or (self.current_epoch % 2 == 1 and optimizer_idx == 1):
        ...
        return loss

in other cases of training_step, it will return None, which will not make any updates to the other optimizer. You might get a warning, but that's fine.
your solution is fine as well, but it will load the data from both the dataloader which is not required in each epoch.

[icedpanda]

Hi @rohitgr7, thanks for the prompt reply.

Just wanna make sure I am on the right track. So I only need to change it to the following and double the epochs

class Datamodule(LightningDataModule):
  def train_dataloader(self):
      if self.current_epoch % 2 == 0:
          A_loader = ...
          return A_loader
      else:
          B_loader = ...
          return B_loader

class Mymodel(LightningModule):
  def training_step(self, batch, batch_idx, optimizer_idx):
      if self.current_epoch % 2 == 0 and optimizer_idx == 0:
           inputs = batch["datasetA"]
           self.model.pretrain()
      elif self.current_epoch % 2 == 1 and optimizer_idx == 1:
          inputs = batch["datasetB"]
          self.model.forward()
      else:
            # if self.current_epoch % 2 == 1 and optimizer_idx == 0:
            # if self.current_epoch % 2 ==0 and optimizer_idx == 1:
            return None
      return loss
      

trainer = Trainer(reload_dataloaders_every_epoch=1, epochs=original_epochs * 2)

so under the hood it will do

for epoch in epochs:
    if not epoch % reload_dataloaders_every_n_epochs:
        # return the right dataloader
        train_loader = model.train_dataloader()
    for batch in train_loader:
        for opt in optimizers:
            # only use opt once and in other cases of `training_step` will return None
            model.forward()
        ...

[rohitgr7]

yes.. correct.
although, I am not sure about this:

inputs = batch["datasetB"]

but it depends on how your dataset is being built. I mean if they are not related, no need to have them renamed as datasetA|datasetB.

[icedpanda]

I have different keys and items in two datasets but yea thanks!