157 lines
5.8 KiB
Python
157 lines
5.8 KiB
Python
"""
|
|
Basic data module implementation
|
|
"""
|
|
import torch
|
|
import logging
|
|
|
|
from torch.utils.data import Dataset
|
|
|
|
from ..label_tensor import LabelTensor
|
|
|
|
|
|
class BaseDataset(Dataset):
|
|
"""
|
|
BaseDataset class, which handle initialization and data retrieval
|
|
:var condition_indices: List of indices
|
|
:var device: torch.device
|
|
"""
|
|
|
|
def __new__(cls, problem=None, device=torch.device('cpu')):
|
|
"""
|
|
Ensure correct definition of __slots__ before initialization
|
|
:param AbstractProblem problem: The formulation of the problem.
|
|
:param torch.device device: The device on which the
|
|
dataset will be loaded.
|
|
"""
|
|
if cls is BaseDataset:
|
|
raise TypeError(
|
|
'BaseDataset cannot be instantiated directly. Use a subclass.')
|
|
if not hasattr(cls, '__slots__'):
|
|
raise TypeError(
|
|
'Something is wrong, __slots__ must be defined in subclasses.')
|
|
return object.__new__(cls)
|
|
|
|
def __init__(self, problem=None, device=torch.device('cpu')):
|
|
""""
|
|
Initialize the object based on __slots__
|
|
:param AbstractProblem problem: The formulation of the problem.
|
|
:param torch.device device: The device on which the
|
|
dataset will be loaded.
|
|
"""
|
|
super().__init__()
|
|
self.empty = True
|
|
self.problem = problem
|
|
self.device = device
|
|
self.condition_indices = None
|
|
for slot in self.__slots__:
|
|
setattr(self, slot, [])
|
|
self.num_el_per_condition = []
|
|
self.conditions_idx = []
|
|
if self.problem is not None:
|
|
self._init_from_problem(self.problem.collector.data_collections)
|
|
self.initialized = False
|
|
|
|
def _init_from_problem(self, collector_dict):
|
|
"""
|
|
TODO
|
|
"""
|
|
for name, data in collector_dict.items():
|
|
keys = list(data.keys())
|
|
if set(self.__slots__) == set(keys):
|
|
self._populate_init_list(data)
|
|
idx = [key for key, val in
|
|
self.problem.collector.conditions_name.items() if
|
|
val == name]
|
|
self.conditions_idx.append(idx)
|
|
self.initialize()
|
|
|
|
def add_points(self, data_dict, condition_idx, batching_dim=0):
|
|
"""
|
|
This method filled internal lists of data points
|
|
:param data_dict: dictionary containing data points
|
|
:param condition_idx: index of the condition to which the data points
|
|
belong to
|
|
:param batching_dim: dimension of the batching
|
|
:raises: ValueError if the dataset has already been initialized
|
|
"""
|
|
if not self.initialized:
|
|
self._populate_init_list(data_dict, batching_dim)
|
|
self.conditions_idx.append(condition_idx)
|
|
self.empty = False
|
|
else:
|
|
raise ValueError('Dataset already initialized')
|
|
|
|
def _populate_init_list(self, data_dict, batching_dim=0):
|
|
current_cond_num_el = None
|
|
for slot in data_dict.keys():
|
|
slot_data = data_dict[slot]
|
|
if batching_dim != 0:
|
|
if isinstance(slot_data, (LabelTensor, torch.Tensor)):
|
|
dims = len(slot_data.size())
|
|
slot_data = slot_data.permute(
|
|
[batching_dim] + [dim for dim in range(dims) if
|
|
dim != batching_dim])
|
|
if current_cond_num_el is None:
|
|
current_cond_num_el = len(slot_data)
|
|
elif current_cond_num_el != len(slot_data):
|
|
raise ValueError('Different dimension in same condition')
|
|
current_list = getattr(self, slot)
|
|
current_list += [slot_data] if not (
|
|
isinstance(slot_data, list)) else slot_data
|
|
self.num_el_per_condition.append(current_cond_num_el)
|
|
|
|
def initialize(self):
|
|
"""
|
|
Initialize the datasets tensors/LabelTensors/lists given the lists
|
|
already filled
|
|
"""
|
|
logging.debug(f'Initialize dataset {self.__class__.__name__}')
|
|
|
|
if self.num_el_per_condition:
|
|
self.condition_indices = torch.cat(
|
|
[
|
|
torch.tensor([i] * self.num_el_per_condition[i],
|
|
dtype=torch.uint8)
|
|
for i in range(len(self.num_el_per_condition))
|
|
],
|
|
dim=0
|
|
)
|
|
for slot in self.__slots__:
|
|
current_attribute = getattr(self, slot)
|
|
if all(isinstance(a, LabelTensor) for a in current_attribute):
|
|
setattr(self, slot, LabelTensor.vstack(current_attribute))
|
|
self.initialized = True
|
|
|
|
def __len__(self):
|
|
"""
|
|
:return: Number of elements in the dataset
|
|
"""
|
|
return len(getattr(self, self.__slots__[0]))
|
|
|
|
def __getitem__(self, idx):
|
|
"""
|
|
:param idx:
|
|
:return:
|
|
"""
|
|
if not isinstance(idx, (tuple, list, slice, int)):
|
|
raise IndexError("Invalid index")
|
|
tensors = []
|
|
for attribute in self.__slots__:
|
|
tensor = getattr(self, attribute)
|
|
if isinstance(attribute, (LabelTensor, torch.Tensor)):
|
|
tensors.append(tensor.__getitem__(idx))
|
|
elif isinstance(attribute, list):
|
|
if isinstance(idx, (list, tuple)):
|
|
tensor = [tensor[i] for i in idx]
|
|
tensors.append(tensor)
|
|
return tensors
|
|
|
|
def apply_shuffle(self, indices):
|
|
for slot in self.__slots__:
|
|
if slot != 'equation':
|
|
attribute = getattr(self, slot)
|
|
if isinstance(attribute, (LabelTensor, torch.Tensor)):
|
|
setattr(self, 'slot', attribute[[indices]])
|
|
if isinstance(attribute, list):
|
|
setattr(self, 'slot', [attribute[i] for i in indices])
|