Support for IterableDatasets

[miguelvr]

🚀 Feature

Pytorch introduced the IterableDataset in v1.2 to allow users to process streams of information. ClassyVision currently only supports map style datasets, it would be nice to extend the support to IterableDatasets, given that they are especially useful to process video streams.

Motivation / Pitch

Map style datasets assume each sample can be read completely independently from each other. In some situations, such as processing video streams, it is extremely expensive to open and close a video stream N times to read N frames. IterableDatasets allow streams to be open once and then results are yielded as requested by the data loader, which is substantially more efficient.

Additional context

There are a few important differences between Map style datasets and iterable datasets, that break the current classy vision dataset paradigm:

1. In an iterable dataset, there is no __getitem__ method, this is replaced by the __iter__ method
2. The __len__ method is optional in iterable dataset
3. Samplers do not work with IterableDatasets, sampling and shuffling has to be handled in the dataset

I've come up with a template dataset called ChunkDataset hides some of this complexity away, which might be nice to help beginner users to get started. Nevertheless, in order to get this working with my code, I had to subclass ClassificationTask to modify it and had to create and entirely new base class for this style of dataset (ClassyDataset is not compatible).

chunk.py

class ChunkDataset(IterableDataset):
    def __init__(self, indices: List, process_fn: Callable):
        """
        Subset of IterableDataset to serve as a base class to process streams of data,
        such as audio, video or text.

        Args:
            indices (list): list of arguments to provide to process_fn
            process_fn (callable): function that processes the indices and returns an iterator
        """
        self.idxs = indices
        self._process_fn = process_fn
        self.epoch = 0
        self.shuffle = False

        # replacement for distributed sampler
        distributed = dist.is_available() and dist.is_initialized()
        if distributed:
            num_replicas = dist.get_world_size()
            rank = dist.get_rank()
            self.idxs = self.idxs[rank::num_replicas]

    def __iter__(self):
        # deterministically shuffle based on epoch
        g = torch.Generator()
        g.manual_seed(self.epoch)
        if self.shuffle:
            indices = torch.randperm(len(self.idxs), generator=g).tolist()
            idxs = [self.idxs[i] for i in indices]
        else:
            idxs = self.idxs

        return self._process_fn(idxs)

    @staticmethod
    def worker_init_fn(worker_id):
        worker_info = torch.utils.data.get_worker_info()
        dataset = worker_info.dataset  # the dataset copy in this worker process
        n_workers = worker_info.num_workers
        dataset.idxs = dataset.idxs[worker_id::n_workers]

    def set_epoch(self, epoch):
        self.epoch = epoch
        return self

    def set_shuffle(self, shuffle: bool = True):
        self.shuffle = shuffle
        return self

    def __len__(self):
        raise NotImplementedError


class ClassyChunkDataset(IterableDataset):
    """
    Class representing a dataset abstraction to wrap a ChunkDataset.

    This class wraps a :class:`ChunkDataset` via the `dataset` attribute
    and configures the dataloaders needed to access the datasets.
    Transforms which need to be applied to the data should be specified in this class.
    ClassyChunkDataset can be instantiated from a configuration file as well.
    """

    def __init__(
        self,
        dataset: ChunkDataset,
        batchsize_per_replica: int,
        shuffle: bool,
        transform: Optional[Union[ClassyTransform, Callable]],
    ) -> None:
        """
        Constructor for a ClassyDataset.

        Args:
            batchsize_per_replica: Positive integer indicating batch size for each
                replica
            shuffle: Whether to shuffle between epochs
            transform: When set, transform to be applied to each sample
            num_samples: When set, this restricts the number of samples provided by
                the dataset
        """
        # Asserts:
        assert is_pos_int(
            batchsize_per_replica
        ), "batchsize_per_replica must be a positive int"
        assert isinstance(shuffle, bool), "shuffle must be a boolean"

        # Assignments:
        self.batchsize_per_replica = batchsize_per_replica
        self.shuffle = shuffle
        self.transform = transform
        self.dataset = dataset

        if self.shuffle:
            self.dataset = self.dataset.set_shuffle()

    @classmethod
    def from_config(cls, config: Dict[str, Any]) -> "ClassyDataset":
        """Instantiates a ClassyDataset from a configuration.

        Args:
            config: A configuration for the ClassyDataset.

        Returns:
            A ClassyDataset instance.
        """
        raise NotImplementedError

    @classmethod
    def parse_config(cls, config: Dict[str, Any]):
        """
        This function parses out common config options.

        Args:
            config: A dict with the following string keys -

                | *batchsize_per_replica* (int): Must be a positive int, batch size
                |    for each replica
                | *use_shuffle* (bool): Whether to enable shuffling for the dataset
                | *num_samples* (int, optional): When set, restricts the number of
                     samples in a dataset
                | *transforms*: list of tranform configurations to be applied in order

        Returns:
            A tuple containing the following variables -
                | *transform_config*: Config for the dataset transform. Can be passed to
                |    :func:`transforms.build_transform`
                | *batchsize_per_replica*: Batch size per replica
                | *shuffle*: Whether we should shuffle between epochs
                | *num_samples*: When set, restricts the number of samples in a dataset
        """
        batchsize_per_replica = config.get("batchsize_per_replica")
        shuffle = config.get("use_shuffle")
        num_samples = config.get("num_samples")
        transform_config = config.get("transforms")
        return transform_config, batchsize_per_replica, shuffle, num_samples

    def __iter__(self):
        for sample in self.dataset:
            if self.transform is not None:
                sample = self.transform(sample)
            yield sample

    def __len__(self):
        return len(self.dataset)

    def iterator(self, *args, **kwargs):
        """
        Returns an iterable which can be used to iterate over the data.

        Args:
            shuffle_seed (int, optional): Seed for the shuffle
            current_phase_id (int, optional): The epoch being fetched. Needed so that
                each epoch has a different shuffle order
        Returns:
            An iterable over the data
        """
        # TODO: Fix naming to be consistent (i.e. everyone uses epoch)
        epoch = kwargs.get("current_phase_id", 0)
        assert isinstance(epoch, int), "Epoch must be an int"

        self.dataset = self.dataset.set_epoch(epoch)

        return DataLoader(
            self,
            batch_size=self.batchsize_per_replica,
            num_workers=kwargs.get("num_workers", 0),
            pin_memory=kwargs.get("pin_memory", False),
            multiprocessing_context=kwargs.get("multiprocessing_context", None),
            worker_init_fn=self.worker_init_fn
        )

    def get_batchsize_per_replica(self):
        """
        Get the batch size per replica.

        Returns:
            The batch size for each replica.
        """
        return self.batchsize_per_replica

    def get_global_batchsize(self):
        """
        Get the global batch size, combined over all the replicas.

        Returns:
            The overall batch size of the dataset.
        """
        return self.get_batchsize_per_replica() * get_world_size()

    @staticmethod
    def worker_init_fn(worker_id):
        worker_info = torch.utils.data.get_worker_info()
        dataset = worker_info.dataset.dataset  # the dataset copy in this worker process
        n_workers = worker_info.num_workers
        dataset.idxs = dataset.idxs[worker_id::n_workers]

[vreis]

What exactly makes ClassyDataset incompatible with IterableDataset? Within FB we do use IterableDataset with Classy Vision without a custom task subclass.

The one issue that comes to mind is that ClassificationTask assumes a __len__ implementation, and we have a tiny dataloader wrapper to provide that. If that's the only thing missing I'll open-source it -- let me know!

[miguelvr]

@vreis ClassyDataset has a default __getitem__ method, I think that would not be compatible with a normal DataLoader, but I might be wrong.

I would appreciate the DataLoader wrapper for the length. That was actually the part that I had to change in my custom task: getting the number of batches from the DataLoader

[miguelvr]

I'm also wondering how do you handle the dataset distribution among different workers. I've done that in the init of my abstraction, but I am wondering if there is a cleaner way

[vreis]

@miguelvr: that __getitem__ is there because .iterator() creates a dataloader passing self. You can override .iterator() and create a dataloader from an IterableDataset instead -- in that case ClassyDataset is simply a wrapper to provide config management. Does that make sense?

I agree that's confusing though -- @aadcock maybe we should rethink this.

[vreis]

For data distribution: let's say your dataset performs a SQL database query. In that case you use an IterableDataset and make each worker generate a different query to select its own subset of data (e.g. row_id % world_size == rank).

Looks like ChunkDataset assumes you have random access to the data (e.g. it all fits in memory) but in that case you can just use a regular Dataset.

[miguelvr]

@miguelvr: that getitem is there because .iterator() creates a dataloader passing self. You can override .iterator() and create a dataloader from an IterableDataset instead -- in that case ClassyDataset is simply a wrapper to provide config management. Does that make sense?

I agree that's confusing though -- @aadcock maybe we should rethink this.

I understand the usage, although I agree it is not super user friendly. I guess it could work if we make a tutorial on the documentation.

Looks like ChunkDataset assumes you have random access to the data (e.g. it all fits in memory) but in that case you can just use a regular Dataset.

@vreis The main purpose of ChunkDataset would be for datasets where each index corresponds to multiple samples, like video frames or video clips

The only thing that still prevents everything from working with ClassificationTask is the len operation over the dataloader (which will have infinite length and raise an error)

[mannatsingh]

@miguelvr I think #455 should solve this issue. We should still document this in a tutorial - I'll create a separate issue for that. Please let me know if you are fine with closing this issue out.

[miguelvr]

indeed, closing the issue