180 lines
6.8 KiB
Python
180 lines
6.8 KiB
Python
"""
|
|
This module provide basic data management functionalities
|
|
"""
|
|
|
|
import math
|
|
import torch
|
|
import logging
|
|
from pytorch_lightning import LightningDataModule
|
|
from .sample_dataset import SamplePointDataset
|
|
from .supervised_dataset import SupervisedDataset
|
|
from .unsupervised_dataset import UnsupervisedDataset
|
|
from .pina_dataloader import PinaDataLoader
|
|
from .pina_subset import PinaSubset
|
|
|
|
|
|
class PinaDataModule(LightningDataModule):
|
|
"""
|
|
This class extend LightningDataModule, allowing proper creation and
|
|
management of different types of Datasets defined in PINA
|
|
"""
|
|
|
|
def __init__(self,
|
|
problem,
|
|
device,
|
|
train_size=.7,
|
|
test_size=.1,
|
|
val_size=.2,
|
|
predict_size=0.,
|
|
batch_size=None,
|
|
shuffle=True,
|
|
datasets=None):
|
|
"""
|
|
Initialize the object, creating dataset based on input problem
|
|
:param AbstractProblem problem: PINA problem
|
|
:param device: Device used for training and testing
|
|
:param train_size: number/percentage of elements in train split
|
|
:param test_size: number/percentage of elements in test split
|
|
:param eval_size: number/percentage of elements in evaluation split
|
|
:param batch_size: batch size used for training
|
|
:param datasets: list of datasets objects
|
|
"""
|
|
logging.debug('Start initialization of Pina DataModule')
|
|
logging.info('Start initialization of Pina DataModule')
|
|
super().__init__()
|
|
self.problem = problem
|
|
self.device = device
|
|
self.dataset_classes = [
|
|
SupervisedDataset, UnsupervisedDataset, SamplePointDataset
|
|
]
|
|
if datasets is None:
|
|
self.datasets = None
|
|
else:
|
|
self.datasets = datasets
|
|
|
|
self.split_length = []
|
|
self.split_names = []
|
|
self.loader_functions = {}
|
|
self.batch_size = batch_size
|
|
self.condition_names = problem.collector.conditions_name
|
|
|
|
if train_size > 0:
|
|
self.split_names.append('train')
|
|
self.split_length.append(train_size)
|
|
self.loader_functions['train_dataloader'] = lambda: PinaDataLoader(
|
|
self.splits['train'], self.batch_size, self.condition_names)
|
|
if test_size > 0:
|
|
self.split_length.append(test_size)
|
|
self.split_names.append('test')
|
|
self.loader_functions['test_dataloader'] = lambda: PinaDataLoader(
|
|
self.splits['test'], self.batch_size, self.condition_names)
|
|
if val_size > 0:
|
|
self.split_length.append(val_size)
|
|
self.split_names.append('val')
|
|
self.loader_functions['val_dataloader'] = lambda: PinaDataLoader(
|
|
self.splits['val'], self.batch_size, self.condition_names)
|
|
if predict_size > 0:
|
|
self.split_length.append(predict_size)
|
|
self.split_names.append('predict')
|
|
self.loader_functions['predict_dataloader'] = lambda: PinaDataLoader(
|
|
self.splits['predict'], self.batch_size, self.condition_names)
|
|
self.splits = {k: {} for k in self.split_names}
|
|
self.shuffle = shuffle
|
|
|
|
for k, v in self.loader_functions.items():
|
|
setattr(self, k, v)
|
|
|
|
def prepare_data(self):
|
|
if self.datasets is None:
|
|
self._create_datasets()
|
|
|
|
def setup(self, stage=None):
|
|
"""
|
|
Perform the splitting of the dataset
|
|
"""
|
|
logging.debug('Start setup of Pina DataModule obj')
|
|
if self.datasets is None:
|
|
self._create_datasets()
|
|
if stage == 'fit' or stage is None:
|
|
for dataset in self.datasets:
|
|
if len(dataset) > 0:
|
|
splits = self.dataset_split(dataset,
|
|
self.split_length,
|
|
shuffle=self.shuffle)
|
|
for i in range(len(self.split_length)):
|
|
self.splits[self.split_names[i]][
|
|
dataset.data_type] = splits[i]
|
|
elif stage == 'test':
|
|
raise NotImplementedError("Testing pipeline not implemented yet")
|
|
else:
|
|
raise ValueError("stage must be either 'fit' or 'test'")
|
|
|
|
@staticmethod
|
|
def dataset_split(dataset, lengths, seed=None, shuffle=True):
|
|
"""
|
|
Perform the splitting of the dataset
|
|
:param dataset: dataset object we wanted to split
|
|
:param lengths: lengths of elements in dataset
|
|
:param seed: random seed
|
|
:param shuffle: shuffle dataset
|
|
:return: split dataset
|
|
:rtype: PinaSubset
|
|
"""
|
|
if sum(lengths) - 1 < 1e-3:
|
|
len_dataset = len(dataset)
|
|
lengths = [
|
|
int(math.floor(len_dataset * length)) for length in lengths
|
|
]
|
|
remainder = len(dataset) - sum(lengths)
|
|
for i in range(remainder):
|
|
lengths[i % len(lengths)] += 1
|
|
elif sum(lengths) - 1 >= 1e-3:
|
|
raise ValueError(f"Sum of lengths is {sum(lengths)} less than 1")
|
|
|
|
if shuffle:
|
|
if seed is not None:
|
|
generator = torch.Generator()
|
|
generator.manual_seed(seed)
|
|
indices = torch.randperm(sum(lengths), generator=generator)
|
|
else:
|
|
indices = torch.randperm(sum(lengths))
|
|
dataset.apply_shuffle(indices)
|
|
|
|
indices = torch.arange(0, sum(lengths), 1, dtype=torch.uint8).tolist()
|
|
offsets = [
|
|
sum(lengths[:i]) if i > 0 else 0 for i in range(len(lengths))
|
|
]
|
|
return [
|
|
PinaSubset(dataset, indices[offset:offset + length])
|
|
for offset, length in zip(offsets, lengths)
|
|
]
|
|
|
|
def _create_datasets(self):
|
|
"""
|
|
Create the dataset objects putting data
|
|
"""
|
|
logging.debug('Dataset creation in PinaDataModule obj')
|
|
collector = self.problem.collector
|
|
batching_dim = self.problem.batching_dimension
|
|
datasets_slots = [i.__slots__ for i in self.dataset_classes]
|
|
self.datasets = [
|
|
dataset(device=self.device) for dataset in self.dataset_classes
|
|
]
|
|
logging.debug('Filling datasets in PinaDataModule obj')
|
|
for name, data in collector.data_collections.items():
|
|
keys = list(data.keys())
|
|
idx = [
|
|
key for key, val in collector.conditions_name.items()
|
|
if val == name
|
|
]
|
|
for i, slot in enumerate(datasets_slots):
|
|
if slot == keys:
|
|
self.datasets[i].add_points(data, idx[0], batching_dim)
|
|
continue
|
|
datasets = []
|
|
for dataset in self.datasets:
|
|
if not dataset.empty:
|
|
dataset.initialize()
|
|
datasets.append(dataset)
|
|
self.datasets = datasets
|