Multiple dataloaders in training

[YuShen1116]

Hi, I'm trying to use integrate pytorch lightning into my current pipeline. But I'm having some difficulties in using multiple dataloaders. In my current use case, let's say I have 10 dataloaders in my pipeline, but for each training step, I'm only sampling data from 5 of them. Is it doable in pytorch lightning? I can sample from 10 dataloaders each step but that would be a waste to system IO and GPU memory. Thanks for any help!

[tchaton]

Hey @YuShen1116,

Here is the pseudo code to get it working

from typing import List
import itertools
from pytorch_lightning.trainer.supporters import CombinedLoader
from pytorch_lightning.utilities.apply_func import apply_to_collection


class CyclingLoader(object):

    def __init__(self, combined_loaders: List[CombinedLoader]):
        self.combined_loaders = combined_loaders
        self._dataloader_idx_cycle = itertools.cycle(range(len(combined_loaders)))

    def __iter__(self):
        self._iterators = apply_to_collection(self.combined_loaders, CombinedLoader, iter)
        self._dataloader_idx_cycle_iter = iter(self._dataloader_idx_cycle)
        return self

    def __next__(self):
        iterator_idx = next(self._dataloader_idx_cycle_iter)
        return next(self._iterators[iterator_idx])


class MyDataModule(DataModule):


    def train_dataloader(self):
        ds_1, .... ds_10 = create_dataloaders()

        ds_1_5 = CombinedLoader([ds_1, ... ds_5])
        ds_6_10 = CombinedLoader([ds_6, ... ds_10])

        return CyclingLoader([ds_1_5, ds_6_10])


class Model(LightningModule):

    def training_step(self, batch, batch_idx):
        if batch_idx % 2 == 0:
            # batches from dataloaders from 1 - 5
        elif batch_idx % 2 == 1:
            # batches from dataloaders from 6 - 10

[YuShen1116]

Thank you so much for your help!