357 lines
14 KiB
Python
357 lines
14 KiB
Python
import bisect
|
|
import warnings
|
|
import math
|
|
from typing import (
|
|
Generic,
|
|
Iterable,
|
|
Iterator,
|
|
List,
|
|
Optional,
|
|
Sequence,
|
|
Tuple,
|
|
TypeVar,
|
|
Union
|
|
)
|
|
|
|
# No 'default_generator' in torch/__init__.pyi
|
|
from torch import default_generator, randperm
|
|
from torch._utils import _accumulate
|
|
|
|
from ... import Generator, Tensor
|
|
|
|
__all__ = [
|
|
"Dataset",
|
|
"IterableDataset",
|
|
"TensorDataset",
|
|
"ConcatDataset",
|
|
"ChainDataset",
|
|
"Subset",
|
|
"random_split",
|
|
]
|
|
|
|
T_co = TypeVar('T_co', covariant=True)
|
|
T = TypeVar('T')
|
|
|
|
|
|
class Dataset(Generic[T_co]):
|
|
r"""An abstract class representing a :class:`Dataset`.
|
|
|
|
All datasets that represent a map from keys to data samples should subclass
|
|
it. All subclasses should overwrite :meth:`__getitem__`, supporting fetching a
|
|
data sample for a given key. Subclasses could also optionally overwrite
|
|
:meth:`__len__`, which is expected to return the size of the dataset by many
|
|
:class:`~torch.utils.data.Sampler` implementations and the default options
|
|
of :class:`~torch.utils.data.DataLoader`.
|
|
|
|
.. note::
|
|
:class:`~torch.utils.data.DataLoader` by default constructs a index
|
|
sampler that yields integral indices. To make it work with a map-style
|
|
dataset with non-integral indices/keys, a custom sampler must be provided.
|
|
"""
|
|
|
|
def __getitem__(self, index) -> T_co:
|
|
raise NotImplementedError
|
|
|
|
def __add__(self, other: 'Dataset[T_co]') -> 'ConcatDataset[T_co]':
|
|
return ConcatDataset([self, other])
|
|
|
|
# No `def __len__(self)` default?
|
|
# See NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ]
|
|
# in pytorch/torch/utils/data/sampler.py
|
|
|
|
|
|
class IterableDataset(Dataset[T_co]):
|
|
r"""An iterable Dataset.
|
|
|
|
All datasets that represent an iterable of data samples should subclass it.
|
|
Such form of datasets is particularly useful when data come from a stream.
|
|
|
|
All subclasses should overwrite :meth:`__iter__`, which would return an
|
|
iterator of samples in this dataset.
|
|
|
|
When a subclass is used with :class:`~torch.utils.data.DataLoader`, each
|
|
item in the dataset will be yielded from the :class:`~torch.utils.data.DataLoader`
|
|
iterator. When :attr:`num_workers > 0`, each worker process will have a
|
|
different copy of the dataset object, so it is often desired to configure
|
|
each copy independently to avoid having duplicate data returned from the
|
|
workers. :func:`~torch.utils.data.get_worker_info`, when called in a worker
|
|
process, returns information about the worker. It can be used in either the
|
|
dataset's :meth:`__iter__` method or the :class:`~torch.utils.data.DataLoader` 's
|
|
:attr:`worker_init_fn` option to modify each copy's behavior.
|
|
|
|
Example 1: splitting workload across all workers in :meth:`__iter__`::
|
|
|
|
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_DATALOADER)
|
|
>>> # xdoctest: +SKIP("Fails on MacOS12")
|
|
>>> class MyIterableDataset(torch.utils.data.IterableDataset):
|
|
... def __init__(self, start, end):
|
|
... super(MyIterableDataset).__init__()
|
|
... assert end > start, "this example code only works with end >= start"
|
|
... self.start = start
|
|
... self.end = end
|
|
...
|
|
... def __iter__(self):
|
|
... worker_info = torch.utils.data.get_worker_info()
|
|
... if worker_info is None: # single-process data loading, return the full iterator
|
|
... iter_start = self.start
|
|
... iter_end = self.end
|
|
... else: # in a worker process
|
|
... # split workload
|
|
... per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
|
|
... worker_id = worker_info.id
|
|
... iter_start = self.start + worker_id * per_worker
|
|
... iter_end = min(iter_start + per_worker, self.end)
|
|
... return iter(range(iter_start, iter_end))
|
|
...
|
|
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
|
|
>>> ds = MyIterableDataset(start=3, end=7)
|
|
|
|
>>> # Single-process loading
|
|
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
|
|
[tensor([3]), tensor([4]), tensor([5]), tensor([6])]
|
|
|
|
>>> # xdoctest: +REQUIRES(POSIX)
|
|
>>> # Mult-process loading with two worker processes
|
|
>>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6].
|
|
>>> # xdoctest: +IGNORE_WANT("non deterministic")
|
|
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
|
|
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]
|
|
|
|
>>> # With even more workers
|
|
>>> # xdoctest: +IGNORE_WANT("non deterministic")
|
|
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12)))
|
|
[tensor([3]), tensor([5]), tensor([4]), tensor([6])]
|
|
|
|
Example 2: splitting workload across all workers using :attr:`worker_init_fn`::
|
|
|
|
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_DATALOADER)
|
|
>>> class MyIterableDataset(torch.utils.data.IterableDataset):
|
|
... def __init__(self, start, end):
|
|
... super(MyIterableDataset).__init__()
|
|
... assert end > start, "this example code only works with end >= start"
|
|
... self.start = start
|
|
... self.end = end
|
|
...
|
|
... def __iter__(self):
|
|
... return iter(range(self.start, self.end))
|
|
...
|
|
>>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6].
|
|
>>> ds = MyIterableDataset(start=3, end=7)
|
|
|
|
>>> # Single-process loading
|
|
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=0)))
|
|
[3, 4, 5, 6]
|
|
>>>
|
|
>>> # Directly doing multi-process loading yields duplicate data
|
|
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2)))
|
|
[3, 3, 4, 4, 5, 5, 6, 6]
|
|
|
|
>>> # Define a `worker_init_fn` that configures each dataset copy differently
|
|
>>> 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
|
|
... overall_start = dataset.start
|
|
... overall_end = dataset.end
|
|
... # configure the dataset to only process the split workload
|
|
... per_worker = int(math.ceil((overall_end - overall_start) / float(worker_info.num_workers)))
|
|
... worker_id = worker_info.id
|
|
... dataset.start = overall_start + worker_id * per_worker
|
|
... dataset.end = min(dataset.start + per_worker, overall_end)
|
|
...
|
|
|
|
>>> # Mult-process loading with the custom `worker_init_fn`
|
|
>>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6].
|
|
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=2, worker_init_fn=worker_init_fn)))
|
|
[3, 5, 4, 6]
|
|
|
|
>>> # With even more workers
|
|
>>> print(list(torch.utils.data.DataLoader(ds, num_workers=12, worker_init_fn=worker_init_fn)))
|
|
[3, 4, 5, 6]
|
|
"""
|
|
def __iter__(self) -> Iterator[T_co]:
|
|
raise NotImplementedError
|
|
|
|
def __add__(self, other: Dataset[T_co]):
|
|
return ChainDataset([self, other])
|
|
|
|
# No `def __len__(self)` default? Subclasses raise `TypeError` when needed.
|
|
# See NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ]
|
|
|
|
|
|
class TensorDataset(Dataset[Tuple[Tensor, ...]]):
|
|
r"""Dataset wrapping tensors.
|
|
|
|
Each sample will be retrieved by indexing tensors along the first dimension.
|
|
|
|
Args:
|
|
*tensors (Tensor): tensors that have the same size of the first dimension.
|
|
"""
|
|
tensors: Tuple[Tensor, ...]
|
|
|
|
def __init__(self, *tensors: Tensor) -> None:
|
|
assert all(tensors[0].size(0) == tensor.size(0) for tensor in tensors), "Size mismatch between tensors"
|
|
self.tensors = tensors
|
|
|
|
def __getitem__(self, index):
|
|
return tuple(tensor[index] for tensor in self.tensors)
|
|
|
|
def __len__(self):
|
|
return self.tensors[0].size(0)
|
|
|
|
|
|
class ConcatDataset(Dataset[T_co]):
|
|
r"""Dataset as a concatenation of multiple datasets.
|
|
|
|
This class is useful to assemble different existing datasets.
|
|
|
|
Args:
|
|
datasets (sequence): List of datasets to be concatenated
|
|
"""
|
|
datasets: List[Dataset[T_co]]
|
|
cumulative_sizes: List[int]
|
|
|
|
@staticmethod
|
|
def cumsum(sequence):
|
|
r, s = [], 0
|
|
for e in sequence:
|
|
l = len(e)
|
|
r.append(l + s)
|
|
s += l
|
|
return r
|
|
|
|
def __init__(self, datasets: Iterable[Dataset]) -> None:
|
|
super().__init__()
|
|
self.datasets = list(datasets)
|
|
assert len(self.datasets) > 0, 'datasets should not be an empty iterable' # type: ignore[arg-type]
|
|
for d in self.datasets:
|
|
assert not isinstance(d, IterableDataset), "ConcatDataset does not support IterableDataset"
|
|
self.cumulative_sizes = self.cumsum(self.datasets)
|
|
|
|
def __len__(self):
|
|
return self.cumulative_sizes[-1]
|
|
|
|
def __getitem__(self, idx):
|
|
if idx < 0:
|
|
if -idx > len(self):
|
|
raise ValueError("absolute value of index should not exceed dataset length")
|
|
idx = len(self) + idx
|
|
dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
|
|
if dataset_idx == 0:
|
|
sample_idx = idx
|
|
else:
|
|
sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
|
|
return self.datasets[dataset_idx][sample_idx]
|
|
|
|
@property
|
|
def cummulative_sizes(self):
|
|
warnings.warn("cummulative_sizes attribute is renamed to "
|
|
"cumulative_sizes", DeprecationWarning, stacklevel=2)
|
|
return self.cumulative_sizes
|
|
|
|
|
|
class ChainDataset(IterableDataset):
|
|
r"""Dataset for chaining multiple :class:`IterableDataset` s.
|
|
|
|
This class is useful to assemble different existing dataset streams. The
|
|
chaining operation is done on-the-fly, so concatenating large-scale
|
|
datasets with this class will be efficient.
|
|
|
|
Args:
|
|
datasets (iterable of IterableDataset): datasets to be chained together
|
|
"""
|
|
def __init__(self, datasets: Iterable[Dataset]) -> None:
|
|
super().__init__()
|
|
self.datasets = datasets
|
|
|
|
def __iter__(self):
|
|
for d in self.datasets:
|
|
assert isinstance(d, IterableDataset), "ChainDataset only supports IterableDataset"
|
|
for x in d:
|
|
yield x
|
|
|
|
def __len__(self):
|
|
total = 0
|
|
for d in self.datasets:
|
|
assert isinstance(d, IterableDataset), "ChainDataset only supports IterableDataset"
|
|
total += len(d) # type: ignore[arg-type]
|
|
return total
|
|
|
|
|
|
class Subset(Dataset[T_co]):
|
|
r"""
|
|
Subset of a dataset at specified indices.
|
|
|
|
Args:
|
|
dataset (Dataset): The whole Dataset
|
|
indices (sequence): Indices in the whole set selected for subset
|
|
"""
|
|
dataset: Dataset[T_co]
|
|
indices: Sequence[int]
|
|
|
|
def __init__(self, dataset: Dataset[T_co], indices: Sequence[int]) -> None:
|
|
self.dataset = dataset
|
|
self.indices = indices
|
|
|
|
def __getitem__(self, idx):
|
|
if isinstance(idx, list):
|
|
return self.dataset[[self.indices[i] for i in idx]]
|
|
return self.dataset[self.indices[idx]]
|
|
|
|
def __len__(self):
|
|
return len(self.indices)
|
|
|
|
|
|
def random_split(dataset: Dataset[T], lengths: Sequence[Union[int, float]],
|
|
generator: Optional[Generator] = default_generator) -> List[Subset[T]]:
|
|
r"""
|
|
Randomly split a dataset into non-overlapping new datasets of given lengths.
|
|
|
|
If a list of fractions that sum up to 1 is given,
|
|
the lengths will be computed automatically as
|
|
floor(frac * len(dataset)) for each fraction provided.
|
|
|
|
After computing the lengths, if there are any remainders, 1 count will be
|
|
distributed in round-robin fashion to the lengths
|
|
until there are no remainders left.
|
|
|
|
Optionally fix the generator for reproducible results, e.g.:
|
|
|
|
Example:
|
|
>>> # xdoctest: +SKIP
|
|
>>> generator1 = torch.Generator().manual_seed(42)
|
|
>>> generator2 = torch.Generator().manual_seed(42)
|
|
>>> random_split(range(10), [3, 7], generator=generator1)
|
|
>>> random_split(range(30), [0.3, 0.3, 0.4], generator=generator2)
|
|
|
|
Args:
|
|
dataset (Dataset): Dataset to be split
|
|
lengths (sequence): lengths or fractions of splits to be produced
|
|
generator (Generator): Generator used for the random permutation.
|
|
"""
|
|
if math.isclose(sum(lengths), 1) and sum(lengths) <= 1:
|
|
subset_lengths: List[int] = []
|
|
for i, frac in enumerate(lengths):
|
|
if frac < 0 or frac > 1:
|
|
raise ValueError(f"Fraction at index {i} is not between 0 and 1")
|
|
n_items_in_split = int(
|
|
math.floor(len(dataset) * frac) # type: ignore[arg-type]
|
|
)
|
|
subset_lengths.append(n_items_in_split)
|
|
remainder = len(dataset) - sum(subset_lengths) # type: ignore[arg-type]
|
|
# add 1 to all the lengths in round-robin fashion until the remainder is 0
|
|
for i in range(remainder):
|
|
idx_to_add_at = i % len(subset_lengths)
|
|
subset_lengths[idx_to_add_at] += 1
|
|
lengths = subset_lengths
|
|
for i, length in enumerate(lengths):
|
|
if length == 0:
|
|
warnings.warn(f"Length of split at index {i} is 0. "
|
|
f"This might result in an empty dataset.")
|
|
|
|
# Cannot verify that dataset is Sized
|
|
if sum(lengths) != len(dataset): # type: ignore[arg-type]
|
|
raise ValueError("Sum of input lengths does not equal the length of the input dataset!")
|
|
|
|
indices = randperm(sum(lengths), generator=generator).tolist() # type: ignore[call-overload]
|
|
return [Subset(dataset, indices[offset - length : offset]) for offset, length in zip(_accumulate(lengths), lengths)]
|