|
|
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()] |
|
|
|
