Conditionally have a single or multiple data Loaders.

[prateeky2806]

Hi, I have a use case where I need to have a single validation dataloader or multiple validation dataloader based on a flag. My data module looks something like

class CodeXGlueDataModule(pl.LightningDataModule):
    def __init__(self, args):
        ...

    def train_dataloader(self):
        ...

    def val_dataloader(self):
        loader1 = ...

        if flag:
            loader2 = ...
            loaders = [dataloader_ppl, dataloader_bleu]
            return loaders
        loaders = [dataloader_ppl, ]
        return loaders

and my lightening model looks like

def validation_step(self, batch, batch_idx, dataloader_idx):
    if dataloader_idx == 0:
        ...
        return {"x":x}

    elif dataloader_idx == 1:
        ...
        return {"y":y}    

When I run the code with flag = True, it works fines but when I run it with flag = False, I get an error saying

validation_step() missing 1 required positional argument: 'dataloader_idx'

I understand that a single dataloader is not treated as a list and hence there is no dataloader_idx argument for the validation_step function but is there something I can do to make this work.

Thanks in Advance!

[akihironitta]

@prateeky2806 Making it an optional keyword argument will avoid the error?

-def validation_step(self, batch, batch_idx, dataloader_idx):
-    if dataloader_idx == 0:
+def validation_step(self, batch, batch_idx, dataloader_idx=None):
+    if dataloader_idx == 0 or dataloader_idx is None:
        ...

    elif dataloader_idx == 1:
        ...

[prateeky2806]

Thanks @akihironitta this solves the problem I mentioned but now the validation_epoch_end is also required to be tackle this as outputs will have different length based on the number of dataloaders.

def validation_epoch_end(self, outputs):
       val_loss = th.stack([out["val_loss"] for out in outputs[0]]).mean()
       val_ppl = th.stack([out["val_ppl"] for out in outputs[1]]).mean()

Is there a way to figure out if there are multiple dataloader in this function?

[akihironitta]

@prateeky2806 Someone else in the community/core might have another better idea, but I don't come up with a better way afaik now...

[prateeky2806]

Something like this works.

def validation_epoch_end(self, outputs):
    if not isinstance(outputs[0], list):
        val_loss = th.stack([out["val_loss"] for out in outputs]).mean()
        val_ppl = th.stack([out["val_ppl"] for out in outputs]).mean()
    else:
        val_loss = th.stack([out["val_loss"] for out in outputs[0]]).mean()
        val_ppl = th.stack([out["val_ppl"] for out in outputs[1]]).mean()