| from typing import List |
|
|
| import torch |
| from torch_geometric.loader import DataLoader |
|
|
| from ._utils import _test_condition |
|
|
|
|
| __all__ = ["Dataset"] |
|
|
|
|
| class Dataset: |
|
|
| def __init__(self, forward_path: str, backward_path: str, splits_path: str, |
| test_indices: bool = True): |
| if forward_path is not None: |
| self.dataset_forward = torch.load(forward_path) |
| else: |
| self.dataset_forward = [] |
| |
| self.dataset_backward = torch.load(backward_path) |
|
|
| self.splits = torch.load(splits_path) |
| if 1 in self.splits.keys(): |
| self.num_of_folds = len(self.splits) |
| else: |
| self.num_of_folds = 1 |
|
|
| |
| if test_indices: |
| if self.num_of_folds > 1: |
| for idx in range(1, self.num_of_folds+ 1): |
| self._test_indices(self.splits[idx]) |
| else: |
| self._test_indices(self.splits) |
|
|
| self.prepare_fold() |
|
|
| def prepare_fold(self, fold_idx: int = 1): |
| if self.num_of_folds > 1: |
| _test_condition(isinstance(fold_idx, int), "'fold_idx' must be an integer") |
| _test_condition(1 <= fold_idx < self.num_of_folds + 1, f"'fold_idx' must be between 1 and {self.num_of_folds+1}") |
| fold_splits = self.splits[fold_idx] |
| else: |
| fold_splits = self.splits |
|
|
| self.train_index_forward = fold_splits["train_index_forward"] |
| self.train_index_backward = fold_splits["train_index_backward"] |
| self.val_index_forward = fold_splits["val_index_forward"] |
| self.val_index_backward = fold_splits["val_index_backward"] |
| self.test_index_forward = fold_splits["test_index_forward"] |
| self.test_index_backward = fold_splits["test_index_backward"] |
|
|
| if self.train_index_forward is not None: |
| self.train_dataset_forward = [self.dataset_forward[i] for i in self.train_index_forward] |
| self.val_dataset_forward = [self.dataset_forward[i] for i in self.val_index_forward] |
| self.test_dataset_forward = [self.dataset_forward[i] for i in self.test_index_forward] |
| self.train_dataset_backward = [self.dataset_backward[i] for i in self.train_index_backward] |
| self.val_dataset_backward = [self.dataset_backward[i] for i in self.val_index_backward] |
| self.test_dataset_backward = [self.dataset_backward[i] for i in self.test_index_backward] |
|
|
| def get_num_vars(self) -> int: |
| return self.dataset_backward[0].num_nodes |
|
|
| def _test_indices(self, fold_splits): |
| |
| if fold_splits["train_index_forward"] is not None: |
| set_trif = set(fold_splits["train_index_forward"]) |
| set_vif = set(fold_splits["val_index_forward"]) |
| set_teif = set(fold_splits["test_index_forward"]) |
| _test_condition(not any(x in set_teif for x in set_trif), "Overlap between train and test indices should be zero (forward)") |
| _test_condition(not any(x in set_vif for x in set_trif), "Overlap between train and validation indices should be zero (forward)") |
| _test_condition(not any(x in set_teif for x in set_vif), "Overlap between validation and test indices should be zero (forward)") |
|
|
|
|
| |
| set_trib = set(fold_splits["train_index_backward"]) |
| set_vib = set(fold_splits["val_index_backward"]) |
| set_teib = set(fold_splits["test_index_backward"]) |
| |
| _test_condition(not any(x in set_vib for x in set_trib), "Overlap between train and validation indices should be zero (backward)") |
| _test_condition(not any(x in set_teib for x in set_trib), "Overlap between train and test indices should be zero (backward)") |
| _test_condition(not any(x in set_teib for x in set_vib), "Overlap between validation and test indices should be zero (backward)") |
|
|
| def get_dataloaders(self, batch_size: int = 64, shuffle = True, **kwargs) -> List[DataLoader]: |
| kwargs = {**{"shuffle": shuffle, "drop_last": True}, **kwargs} |
| return [ |
| DataLoader(self.train_dataset_forward, batch_size=batch_size, **kwargs) if hasattr(self, 'train_dataset_forward') else None, |
| DataLoader(self.train_dataset_backward, batch_size=batch_size, **kwargs), |
| DataLoader(self.val_dataset_forward, batch_size=batch_size, **kwargs) if hasattr(self, 'val_dataset_forward') else None, |
| DataLoader(self.val_dataset_backward, batch_size=batch_size, **kwargs), |
| DataLoader(self.test_dataset_forward, batch_size=batch_size, **kwargs) if hasattr(self, 'test_dataset_forward') else None, |
| DataLoader(self.test_dataset_backward, batch_size=batch_size, **kwargs) |
| ] |
|
|
