Multiple dataloaders in training
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!
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