| | from typing import Iterator, Optional |
| |
|
| | import torch |
| | from torch.utils.data import Dataset, IterableDataset |
| |
|
| |
|
| | class ValidationWrapper(Dataset): |
| | """Wraps a dataset so that PyTorch Lightning's validation step can be turned into a |
| | visualization step. |
| | """ |
| |
|
| | dataset: Dataset |
| | dataset_iterator: Optional[Iterator] |
| | length: int |
| |
|
| | def __init__(self, dataset: Dataset, length: int) -> None: |
| | super().__init__() |
| | self.dataset = dataset |
| | self.length = length |
| | self.dataset_iterator = None |
| |
|
| | def __len__(self): |
| | return self.length |
| | |
| | def __getitem__(self, index: tuple): |
| | if isinstance(self.dataset, IterableDataset): |
| | if self.dataset_iterator is None: |
| | self.dataset_iterator = iter(self.dataset) |
| | return next(self.dataset_iterator) |
| | |
| | random_index = torch.randint(0, len(self.dataset), tuple()) |
| | random_context_num = torch.randint(2, self.dataset.view_sampler.num_context_views + 1, tuple()) |
| | |
| | return self.dataset[random_index.item(), random_context_num.item()] |
| |
|
