sconce.data_generators.base — Pytorch-Sconce master documentation

# Source code for sconce.data_generators.base

    from sconce.datasets import Subset
    from torch.utils import data

    import copy
    import numpy as np
    import torch

    class DataGenerator:
        """
        A thin wrapper around a :py:class:`~torch.utils.data.DataLoader` that
        automatically yields tuples of :py:class:`torch.Tensor` (that
        live on cpu or on cuda).
        A DataGenerator will iterate endlessly.

        Like the underlying :py:class:`~torch.utils.data.DataLoader`, a
        DataGenerator's ``__next__`` method yields two values, which we refer to as
        the `inputs` and the `targets`.

        Arguments:
            data_loader (:py:class:`~torch.utils.data.DataLoader`): the wrapped
                data_loader.

        Note:
            DataGenerators are deprecated as of 1.2.0 and will be removed soon.  Please use a
            :py:class:`~sconce.data_feeds.base.DataFeed` instead.
        """

        def __init__(self, data_loader):
            print("WARNING: DataGenerators are deprecated as of 1.2.0 and will be removed soon.  "
                  "Please use a DataFeed instead")
            self.data_loader = data_loader
            self._inputs_cuda = False
            self._targets_cuda = False
            self.reset()

        def cuda(self, device=None):
            """
            Put the `inputs` and `targets` (yielded by this DataGenerator) on the
            specified device.

            Arguments:
                device (int or bool or dict): if int or bool, sets the behavior for
                    both `inputs` and `targets`. To set them individually, pass a
                    dictionary with keys {'inputs', 'targets'} instead. See
                    :py:meth:`torch.Tensor.cuda` for details.

            Example:
                >>> g = DataGenerator.from_dataset(dataset, batch_size=100)
                >>> g.cuda()
                >>> g.next()
                (Tensor containing:
                 [torch.cuda.FloatTensor of size 100x1x28x28 (GPU 0)],
                 Tensor containing:
                 [torch.cuda.LongTensor of size 100 (GPU 0)])
                >>> g.cuda(False)
                >>> g.next()
                (Tensor containing:
                 [torch.FloatTensor of size 100x1x28x28],
                 Tensor containing:
                 [torch.LongTensor of size 100])
                >>> g.cuda(device={'inputs':0, 'targets':1})
                >>> g.next()
                (Tensor containing:
                 [torch.cuda.FloatTensor of size 100x1x28x28 (GPU 0)],
                 Tensor containing:
                 [torch.cuda.LongTensor of size 100 (GPU 1)])
            """
            if isinstance(device, dict):
                for key, value in device.items():
                    if key == 'inputs':
                        self._inputs_cuda = value
                    elif key == 'targets':
                        self._targets_cuda = value
                    else:
                        raise RuntimeError(f"Invalid key for 'device' argument: "
                                f"({key}) expected to be in ('inputs', 'targets').")
            else:
                self._inputs_cuda = device
                self._targets_cuda = device

        def reset(self):
            """
            Start iterating through the data_loader from the begining.
            """
            self._iterator = iter(self.data_loader)

        @property
        def dataset(self):
            """
            the wrapped data_loader's :py:class:`~torch.utils.data.Dataset`
            """
            return self.data_loader.dataset

        @property
        def batch_size(self):
            """
            the wrapped data_loader's batch_size
            """
            return self.data_loader.batch_size

        @property
        def num_samples(self):
            """
            the ``len`` of the wrapped data_loader's
            :py:class:`~torch.utils.data.Dataset`
            """
            return len(self.dataset)

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

        def __iter__(self):
            return self

        def __next__(self):
            return self.next()

        def preprocess(self, inputs, targets):
            if self._inputs_cuda is False:
                inputs = inputs.cpu()
            else:
                inputs = inputs.cuda(self._inputs_cuda)

            if self._targets_cuda is False:
                if isinstance(targets, torch.Tensor):
                    targets = targets.cpu()
            else:
                targets = targets.cuda(self._targets_cuda)

            return inputs, targets

        def next(self):
            try:
                inputs, targets = self._iterator.next()
            except StopIteration:
                self.reset()
                inputs, targets = self._iterator.next()
            return self.preprocess(inputs, targets)

        @classmethod
        def from_dataset(cls, dataset, split=None, **kwargs):
            """
            Create a DataGenerator from an instantiated dataset.

            Arguments:
                dataset (:py:class:`~torch.utils.data.Dataset`): the pytorch
                    dataset.
                split (float, optional): If not ``None``, it specifies the fraction of the dataset that should be placed
                    into the first of two data_generators.  The remaining data is used for the second data_generator.  Both
                    data_generators will be returned.
                **kwargs: passed directly to the
                    :py:class:`~torch.utils.data.DataLoader`) constructor.
            """
            if split is not None:
                num_samples = len(dataset)
                indices = np.arange(0, num_samples)
                np.random.shuffle(indices)

                dataset1 = dataset
                dataset2 = copy.copy(dataset)

                subset1 = Subset(dataset1, indices=indices[:int(num_samples * split)])
                subset2 = Subset(dataset2, indices=indices[int(num_samples * split):])

                loader1 = data.DataLoader(subset1, **kwargs)
                loader2 = data.DataLoader(subset2, **kwargs)
                return cls(loader1), cls(loader2)
            else:
                data_loader = data.DataLoader(dataset, **kwargs)
                return cls(data_loader)