|
|
from torch.utils.data.datapipes._decorator import functional_datapipe |
|
|
from torch.utils.data.datapipes.datapipe import MapDataPipe |
|
|
from typing import Sized, Tuple, TypeVar |
|
|
|
|
|
__all__ = ["ConcaterMapDataPipe", "ZipperMapDataPipe"] |
|
|
|
|
|
T_co = TypeVar('T_co', covariant=True) |
|
|
|
|
|
|
|
|
@functional_datapipe('concat') |
|
|
class ConcaterMapDataPipe(MapDataPipe): |
|
|
r""" |
|
|
Concatenate multiple Map DataPipes (functional name: ``concat``). |
|
|
The new index of is the cumulative sum of source DataPipes. |
|
|
For example, if there are 2 source DataPipes both with length 5, |
|
|
index 0 to 4 of the resulting `ConcatMapDataPipe` would refer to |
|
|
elements of the first DataPipe, and 5 to 9 would refer to elements |
|
|
of the second DataPipe. |
|
|
|
|
|
Args: |
|
|
datapipes: Map DataPipes being concatenated |
|
|
|
|
|
Example: |
|
|
>>> # xdoctest: +SKIP |
|
|
>>> from torchdata.datapipes.map import SequenceWrapper |
|
|
>>> dp1 = SequenceWrapper(range(3)) |
|
|
>>> dp2 = SequenceWrapper(range(3)) |
|
|
>>> concat_dp = dp1.concat(dp2) |
|
|
>>> list(concat_dp) |
|
|
[0, 1, 2, 0, 1, 2] |
|
|
""" |
|
|
datapipes: Tuple[MapDataPipe] |
|
|
length: int |
|
|
|
|
|
def __init__(self, *datapipes: MapDataPipe): |
|
|
if len(datapipes) == 0: |
|
|
raise ValueError("Expected at least one DataPipe, but got nothing") |
|
|
if not all(isinstance(dp, MapDataPipe) for dp in datapipes): |
|
|
raise TypeError("Expected all inputs to be `MapDataPipe`") |
|
|
if not all(isinstance(dp, Sized) for dp in datapipes): |
|
|
raise TypeError("Expected all inputs to be `Sized`") |
|
|
self.datapipes = datapipes |
|
|
self.length = -1 |
|
|
|
|
|
def __getitem__(self, index) -> T_co: |
|
|
offset = 0 |
|
|
for dp in self.datapipes: |
|
|
if index - offset < len(dp): |
|
|
return dp[index - offset] |
|
|
else: |
|
|
offset += len(dp) |
|
|
raise IndexError("Index {} is out of range.".format(index)) |
|
|
|
|
|
def __len__(self) -> int: |
|
|
if self.length == -1: |
|
|
self.length = sum(len(dp) for dp in self.datapipes) |
|
|
return self.length |
|
|
|
|
|
|
|
|
@functional_datapipe('zip') |
|
|
class ZipperMapDataPipe(MapDataPipe[Tuple[T_co, ...]]): |
|
|
r""" |
|
|
Aggregates elements into a tuple from each of the input DataPipes (functional name: ``zip``). |
|
|
This MataPipe is out of bound as soon as the shortest input DataPipe is exhausted. |
|
|
|
|
|
Args: |
|
|
*datapipes: Map DataPipes being aggregated |
|
|
|
|
|
Example: |
|
|
>>> # xdoctest: +SKIP |
|
|
>>> from torchdata.datapipes.map import SequenceWrapper |
|
|
>>> dp1 = SequenceWrapper(range(3)) |
|
|
>>> dp2 = SequenceWrapper(range(10, 13)) |
|
|
>>> zip_dp = dp1.zip(dp2) |
|
|
>>> list(zip_dp) |
|
|
[(0, 10), (1, 11), (2, 12)] |
|
|
""" |
|
|
datapipes: Tuple[MapDataPipe[T_co], ...] |
|
|
length: int |
|
|
|
|
|
def __init__(self, *datapipes: MapDataPipe[T_co]) -> None: |
|
|
if len(datapipes) == 0: |
|
|
raise ValueError("Expected at least one DataPipe, but got nothing") |
|
|
if not all(isinstance(dp, MapDataPipe) for dp in datapipes): |
|
|
raise TypeError("Expected all inputs to be `MapDataPipe`") |
|
|
if not all(isinstance(dp, Sized) for dp in datapipes): |
|
|
raise TypeError("Expected all inputs to be `Sized`") |
|
|
self.datapipes = datapipes |
|
|
self.length = -1 |
|
|
|
|
|
def __getitem__(self, index) -> Tuple[T_co, ...]: |
|
|
res = [] |
|
|
for dp in self.datapipes: |
|
|
try: |
|
|
res.append(dp[index]) |
|
|
except IndexError: |
|
|
raise IndexError(f"Index {index} is out of range for one of the input MapDataPipes {dp}.") |
|
|
return tuple(res) |
|
|
|
|
|
def __len__(self) -> int: |
|
|
if self.length == -1: |
|
|
self.length = min(len(dp) for dp in self.datapipes) |
|
|
return self.length |
|
|
|
