multimodal/examples/common/data/multidata.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import random
import warnings
from functools import partial
from typing import Callable, List, Optional

import torch
from pytorch_lightning import LightningDataModule


class MultiDataLoader:
    # NOTE: Please check MMF's MultiDataLoader if you want to support
    # epoch based sampling funcs.
    def __init__(
        self,
        loaders: List[torch.utils.data.DataLoader],
        sampling_func: Optional[Callable] = None,
    ):
        """MultiDataLoader takes in a list of dataloaders and a sampling function
        and cycles between these dataloaders after each batch based on the index
        provided by the sampling function passed. Useful for doing multi-tasking
        over multiple datasets

        Args:
            loaders (List[torch.utils.data.DataLoader]): List of dataloaders on
                which the multitasking has to be done.

            sampling_func (Optional[Callable], optional): Function which will return
                the next index to be selected. Defaults to equally weight sampling.
        """
        if loaders is None or len(loaders) == 0:
            warnings.warn(
                "Empty loaders passed into MultiDataLoader. This can have "
                "unintended consequences."
            )

        if sampling_func is None:
            sampling_func = partial(random.choice, range(len(loaders)))

        self.sampling_func = sampling_func
        self.loaders = loaders
        self.num_datasets = len(self.loaders)
        self.iterators = [None for _ in loaders]
        self.current_index = 0
        self.set_samplers()

    def set_samplers(self):
        self.samplers: List[torch.utils.data.Sampler] = []
        for loader in self.loaders:
            if hasattr(loader, "sampler"):
                self.samplers.append(loader.sampler)

    def __iter__(self):
        self.iterators = []

        for loader in self.loaders:
            self.iterators.append(iter(loader))

        self.change_dataloader()

        return self

    def __next__(self):
        """
        Calculation of next batch is performed using following logic.

        Current chosen iterator is set in the change_dataloader function
        based on the `sampling_func` function passed to `__init__` of the
        dataloader which is called to get the index of next selected dataloader.

        If we get the next batch from iterator without any StopIteration exception,
        we return it as it is.

        Epochs don't make sense in case of using `sampling_func` unless you add
        extra logic to support epoch-based sampling functions. MMF does this in
        a different way, so take a look at IterationStrategies there to understand
        how this can be possibly done.

        Think of a case of random (equal) proportional sampling for dataset x and y
        where x is half the size of y. When x will complete its 2 epochs, y will
        have only 1 epoch completed. **So please don't use max_epochs or epoch
        based training in this case as it won't be honored**. If an iterator is
        finished, we just reignite it in this case and finished iterators
        variable isn't used. This means that this case will never reach the
        __iter__ function ever again.


        Returns:
            Dict: Contains two keys, one "batch" containing the batch from current
                selected dataloader and "datamodule_index" which is index of
                currently selected dataloader.
        """
        self.change_dataloader()
        try:
            next_batch = next(self.current_iterator)
        except StopIteration:
            iterator = iter(self.loaders[self.current_index])
            self.iterators[self.current_index] = iterator
            self.current_iterator = iterator
            next_batch = next(self.current_iterator)

        return {"batch": next_batch, "datamodule_index": self.current_index}

    def change_dataloader(self):
        choice = 0

        if self.num_datasets <= 1:
            self.current_index = choice
            self.current_iterator = self.iterators[self.current_index]
            return

        choice = [self.sampling_func()]
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            # This broadcast is probably unnecessary with lightning if everything
            # is already properly seeded. But,to be on safe side, we can still
            # do this.
            # There are also some smarter ways to do this to avoid any broadcasting
            # by basically having a fixed generator with a fixed seed which will
            # always work deterministically.
            # TODO: Check if not doing this provides any speed benefits.
            torch.distributed.broadcast_object_list(choice, 0)

        self.current_index = choice[0]
        self.current_iterator = self.iterators[self.current_index]

    def set_epoch(self, epoch: int):
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            for sampler in self.samplers:
                if sampler is not None and hasattr(sampler, "set_epoch"):
                    sampler.set_epoch(epoch)


class MultiDataModule(LightningDataModule):
    """MultiDataModule is just an abstraction over MultiDataLoader
    that will allow us to integrate it with Lightning.
    """

    # NOTE: Add rest of the functions that should be called on child datamodules
    # as required
    def __init__(
        self,
        datamodules: List[LightningDataModule],
        sampling_func: Optional[Callable] = None,
    ):
        super().__init__()
        self.datamodules = datamodules
        self.sampling_func = sampling_func
        self.current_datamodule_idx = 0

    def setup(self, stage=None):
        for datamodule in self.datamodules:
            datamodule.setup(stage)

    def prepare_data(self):
        for datamodule in self.datamodules:
            datamodule.prepare_data()

    def train_dataloader(self) -> MultiDataLoader:
        # TODO: Fix assign inconsistency
        return self._build_multi_dataloader("train")

    def val_dataloader(self) -> MultiDataLoader:
        return self._build_multi_dataloader("val")

    def test_dataloader(self) -> MultiDataLoader:
        return self._build_multi_dataloader("test")

    def _build_multi_dataloader(self, split="train"):
        dataloaders = []
        for datamodule in self.datamodules:
            dataloaders.append(getattr(datamodule, f"{split}_dataloader")())

        return MultiDataLoader(dataloaders, self.sampling_func)

    def on_before_batch_transfer(self, batch, *args):
        batch, index = batch["batch"], batch["datamodule_index"]
        self.current_datamodule_idx = index
        return self.datamodules[self.current_datamodule_idx].on_before_batch_transfer(
            batch, *args
        )

    def on_after_batch_transfer(self, batch, *args):
        return self.datamodules[self.current_datamodule_idx].on_after_batch_transfer(
            batch, *args
        )

    def teardown(self, stage):
        for datamodule in self.datamodules:
            datamodule.teardown(stage)