| |
| |
| |
| |
|
|
| from typing import Any, Optional, Tuple |
|
|
| from torch.utils.data import Dataset |
|
|
|
|
| def extend_samples_with_index(dataset_class): |
| class DatasetWithIndex(dataset_class): |
| def __init__(self, **kwargs) -> None: |
| root = dataset_class.get_root() |
| super().__init__(root=root, **kwargs) |
|
|
| def __getitem__(self, index: int): |
| image, target = super().__getitem__(index) |
| return image, target, index |
|
|
| return DatasetWithIndex |
|
|
|
|
| class DatasetWithEnumeratedTargets(Dataset): |
| """ |
| If pad_dataset is set, pads based on torch's DistributedSampler implementation, which |
| with drop_last=False pads the last batch to be a multiple of the world size. |
| https://github.com/pytorch/pytorch/blob/main/torch/utils/data/distributed.py#L91 |
| """ |
|
|
| def __init__(self, dataset: Dataset, pad_dataset: bool = False, num_replicas: Optional[int] = None): |
| self._dataset = dataset |
| self._size = len(self._dataset) |
| self._padded_size = self._size |
| self._pad_dataset = pad_dataset |
| if self._pad_dataset: |
| assert num_replicas is not None, "num_replicas should be set if pad_dataset is True" |
| self._padded_size = num_replicas * ((len(dataset) + num_replicas - 1) // num_replicas) |
|
|
| def get_image_relpath(self, index: int) -> str: |
| assert self._pad_dataset or index < self._size |
| return self._dataset.get_image_relpath(index % self._size) |
|
|
| def get_image_data(self, index: int) -> bytes: |
| assert self._pad_dataset or index < self._size |
| return self._dataset.get_image_data(index % self._size) |
|
|
| def get_target(self, index: int) -> Tuple[Any, int]: |
| target = self._dataset.get_target(index % self._size) |
| if index >= self._size: |
| assert self._pad_dataset |
| return (-1, target) |
| return (index, target) |
|
|
| def get_sample_decoder(self, index: int) -> Any: |
| assert self._pad_dataset or index < self._size |
| return self._dataset.get_sample_decoder(index % self._size) |
|
|
| def __getitem__(self, index: int) -> Tuple[Any, Tuple[Any, int]]: |
| image, target = self._dataset[index % self._size] |
| if index >= self._size: |
| assert self._pad_dataset |
| return image, (-1, target) |
| target = index if target is None else target |
| return image, (index, target) |
|
|
| def __len__(self) -> int: |
| return self._padded_size |
|
|
