refact dataset, dataloader and datamodule
This commit is contained in:
FilippoOlivo
@@ -1,9 +1,10 @@
|
||||
"""Module for the PINA dataset classes."""
|
||||
|
||||
from abc import abstractmethod, ABC
|
||||
from torch.utils.data import Dataset
|
||||
from torch_geometric.data import Data
|
||||
from ..graph import Graph, LabelBatch
|
||||
from ..label_tensor import LabelTensor
|
||||
import torch
|
||||
|
||||
|
||||
class PinaDatasetFactory:
|
||||
@@ -41,286 +42,156 @@ class PinaDatasetFactory:
|
||||
if len(conditions_dict) == 0:
|
||||
raise ValueError("No conditions provided")
|
||||
|
||||
dataset_dict = {}
|
||||
|
||||
# Check is a Graph is present in the conditions
|
||||
is_graph = cls._is_graph_dataset(conditions_dict)
|
||||
if is_graph:
|
||||
# If a Graph is present, return a PinaGraphDataset
|
||||
return PinaGraphDataset(conditions_dict, **kwargs)
|
||||
# If no Graph is present, return a PinaTensorDataset
|
||||
return PinaTensorDataset(conditions_dict, **kwargs)
|
||||
for name, data in conditions_dict.items():
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError(
|
||||
f"Condition '{name}' data must be a dictionary"
|
||||
)
|
||||
|
||||
# is_graph = cls._is_graph_dataset(conditions_dict)
|
||||
# if is_graph:
|
||||
# raise NotImplementedError("PinaGraphDataset is not implemented yet.")
|
||||
|
||||
dataset_dict[name] = PinaTensorDataset(data, **kwargs)
|
||||
return dataset_dict
|
||||
|
||||
@staticmethod
|
||||
def _is_graph_dataset(conditions_dict):
|
||||
def _is_graph_dataset(cond_data):
|
||||
"""
|
||||
Check if a graph is present in the conditions (at least one time).
|
||||
|
||||
:param conditions_dict: Dictionary containing the conditions.
|
||||
:type conditions_dict: dict
|
||||
:return: True if a graph is present in the conditions, False otherwise.
|
||||
:rtype: bool
|
||||
TODO: Docstring
|
||||
"""
|
||||
|
||||
# Iterate over the conditions dictionary
|
||||
for v in conditions_dict.values():
|
||||
# Iterate over the values of the current condition
|
||||
for cond in v.values():
|
||||
# Check if the current value is a list of Data objects
|
||||
if isinstance(cond, (Data, Graph, list, tuple)):
|
||||
return True
|
||||
# Iterate over the values of the current condition
|
||||
for cond in cond_data.values():
|
||||
if isinstance(cond, (Data, Graph, list, tuple)):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class PinaDataset(Dataset, ABC):
|
||||
class PinaTensorDataset(Dataset):
|
||||
"""
|
||||
Abstract class for the PINA dataset which extends the PyTorch
|
||||
:class:`~torch.utils.data.Dataset` class. It defines the common interface
|
||||
for :class:`~pina.data.dataset.PinaTensorDataset` and
|
||||
:class:`~pina.data.dataset.PinaGraphDataset` classes.
|
||||
Dataset class for the PINA dataset with :class:`torch.Tensor` and
|
||||
:class:`~pina.label_tensor.LabelTensor` data.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, conditions_dict, max_conditions_lengths, automatic_batching
|
||||
):
|
||||
def __init__(self, data_dict, automatic_batching=None):
|
||||
"""
|
||||
Initialize the instance by storing the conditions dictionary, the
|
||||
maximum number of items per conditions to consider, and the automatic
|
||||
batching flag.
|
||||
Initialize the instance by storing the conditions dictionary.
|
||||
|
||||
:param dict conditions_dict: A dictionary mapping condition names to
|
||||
their respective data. Each key represents a condition name, and the
|
||||
corresponding value is a dictionary containing the associated data.
|
||||
:param dict max_conditions_lengths: Maximum number of data points that
|
||||
can be included in a single batch per condition.
|
||||
:param bool automatic_batching: Indicates whether PyTorch automatic
|
||||
batching is enabled in
|
||||
:class:`~pina.data.data_module.PinaDataModule`.
|
||||
"""
|
||||
|
||||
# Store the conditions dictionary
|
||||
self.conditions_dict = conditions_dict
|
||||
# Store the maximum number of conditions to consider
|
||||
self.max_conditions_lengths = max_conditions_lengths
|
||||
# Store length of each condition
|
||||
self.conditions_length = {
|
||||
k: len(v["input"]) for k, v in self.conditions_dict.items()
|
||||
}
|
||||
# Store the maximum length of the dataset
|
||||
self.length = max(self.conditions_length.values())
|
||||
# Dynamically set the getitem function based on automatic batching
|
||||
if automatic_batching:
|
||||
self._getitem_func = self._getitem_int
|
||||
else:
|
||||
self._getitem_func = self._getitem_dummy
|
||||
|
||||
def _get_max_len(self):
|
||||
"""
|
||||
Returns the length of the longest condition in the dataset.
|
||||
|
||||
:return: Length of the longest condition in the dataset.
|
||||
:rtype: int
|
||||
"""
|
||||
|
||||
max_len = 0
|
||||
for condition in self.conditions_dict.values():
|
||||
max_len = max(max_len, len(condition["input"]))
|
||||
return max_len
|
||||
self.data = data_dict
|
||||
self.automatic_batching = (
|
||||
automatic_batching if automatic_batching is not None else True
|
||||
)
|
||||
self.stack_fn = (
|
||||
{}
|
||||
) # LabelTensor.stack if any(isinstance(v, LabelTensor) for v in data_dict.values()) else torch.stack
|
||||
for k, v in data_dict.items():
|
||||
if isinstance(v, LabelTensor):
|
||||
self.stack_fn[k] = LabelTensor.stack
|
||||
elif isinstance(v, torch.Tensor):
|
||||
self.stack_fn[k] = torch.stack
|
||||
elif isinstance(v, list) and all(
|
||||
isinstance(item, (Data, Graph)) for item in v
|
||||
):
|
||||
self.stack_fn[k] = LabelBatch.from_data_list
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported data type for stacking: {type(v)}"
|
||||
)
|
||||
|
||||
def __len__(self):
|
||||
return self.length
|
||||
return len(next(iter(self.data.values())))
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self._getitem_func(idx)
|
||||
|
||||
def _getitem_dummy(self, idx):
|
||||
"""
|
||||
Return the index itself. This is used when automatic batching is
|
||||
disabled to postpone the data retrieval to the dataloader.
|
||||
|
||||
:param int idx: Index.
|
||||
:return: Index.
|
||||
:rtype: int
|
||||
"""
|
||||
|
||||
# If automatic batching is disabled, return the data at the given index
|
||||
return idx
|
||||
|
||||
def _getitem_int(self, idx):
|
||||
"""
|
||||
Return the data at the given index in the dataset. This is used when
|
||||
automatic batching is enabled.
|
||||
Return the data at the given index in the dataset.
|
||||
|
||||
:param int idx: Index.
|
||||
:return: A dictionary containing the data at the given index.
|
||||
:rtype: dict
|
||||
"""
|
||||
|
||||
# If automatic batching is enabled, return the data at the given index
|
||||
return {
|
||||
k: {k_data: v[k_data][idx % len(v["input"])] for k_data in v.keys()}
|
||||
for k, v in self.conditions_dict.items()
|
||||
}
|
||||
if self.automatic_batching:
|
||||
# Return the data at the given index
|
||||
return {
|
||||
field_name: data[idx] for field_name, data in self.data.items()
|
||||
}
|
||||
return idx
|
||||
|
||||
def get_all_data(self):
|
||||
"""
|
||||
Return all data in the dataset.
|
||||
|
||||
:return: A dictionary containing all the data in the dataset.
|
||||
:rtype: dict
|
||||
"""
|
||||
to_return_dict = {}
|
||||
for condition, data in self.conditions_dict.items():
|
||||
len_condition = len(
|
||||
data["input"]
|
||||
) # Length of the current condition
|
||||
to_return_dict[condition] = self._retrive_data(
|
||||
data, list(range(len_condition))
|
||||
) # Retrieve the data from the current condition
|
||||
return to_return_dict
|
||||
|
||||
def fetch_from_idx_list(self, idx):
|
||||
def _getitem_from_list(self, idx_list):
|
||||
"""
|
||||
Return data from the dataset given a list of indices.
|
||||
|
||||
:param list[int] idx: List of indices.
|
||||
:param list[int] idx_list: List of indices.
|
||||
:return: A dictionary containing the data at the given indices.
|
||||
:rtype: dict
|
||||
"""
|
||||
|
||||
to_return_dict = {}
|
||||
for condition, data in self.conditions_dict.items():
|
||||
# Get the indices for the current condition
|
||||
cond_idx = idx[: self.max_conditions_lengths[condition]]
|
||||
# Get the length of the current condition
|
||||
condition_len = self.conditions_length[condition]
|
||||
# If the length of the dataset is greater than the length of the
|
||||
# current condition, repeat the indices
|
||||
if self.length > condition_len:
|
||||
cond_idx = [idx % condition_len for idx in cond_idx]
|
||||
# Retrieve the data from the current condition
|
||||
to_return_dict[condition] = self._retrive_data(data, cond_idx)
|
||||
return to_return_dict
|
||||
|
||||
@abstractmethod
|
||||
def _retrive_data(self, data, idx_list):
|
||||
"""
|
||||
Abstract method to retrieve data from the dataset given a list of
|
||||
indices.
|
||||
"""
|
||||
|
||||
|
||||
class PinaTensorDataset(PinaDataset):
|
||||
"""
|
||||
Dataset class for the PINA dataset with :class:`torch.Tensor` and
|
||||
:class:`~pina.label_tensor.LabelTensor` data.
|
||||
"""
|
||||
|
||||
# Override _retrive_data method for torch.Tensor data
|
||||
def _retrive_data(self, data, idx_list):
|
||||
"""
|
||||
Retrieve data from the dataset given a list of indices.
|
||||
|
||||
:param dict data: Dictionary containing the data
|
||||
(only :class:`torch.Tensor` or
|
||||
:class:`~pina.label_tensor.LabelTensor`).
|
||||
:param list[int] idx_list: indices to retrieve.
|
||||
:return: Dictionary containing the data at the given indices.
|
||||
:rtype: dict
|
||||
"""
|
||||
|
||||
return {k: v[idx_list] for k, v in data.items()}
|
||||
|
||||
@property
|
||||
def input(self):
|
||||
"""
|
||||
Return the input data for the dataset.
|
||||
|
||||
:return: Dictionary containing the input points.
|
||||
:rtype: dict
|
||||
"""
|
||||
return {k: v["input"] for k, v in self.conditions_dict.items()}
|
||||
|
||||
def update_data(self, new_conditions_dict):
|
||||
"""
|
||||
Update the dataset with new data.
|
||||
This method is used to update the dataset with new data. It replaces
|
||||
the current data with the new data provided in the new_conditions_dict
|
||||
parameter.
|
||||
|
||||
:param dict new_conditions_dict: Dictionary containing the new data.
|
||||
:return: None
|
||||
"""
|
||||
for condition, data in new_conditions_dict.items():
|
||||
if condition in self.conditions_dict:
|
||||
self.conditions_dict[condition].update(data)
|
||||
to_return = {}
|
||||
for field_name, data in self.data.items():
|
||||
if self.stack_fn[field_name] == LabelBatch.from_data_list:
|
||||
to_return[field_name] = self.stack_fn[field_name](
|
||||
[data[i] for i in idx_list]
|
||||
)
|
||||
else:
|
||||
self.conditions_dict[condition] = data
|
||||
to_return[field_name] = data[idx_list]
|
||||
return to_return
|
||||
|
||||
|
||||
class PinaGraphDataset(PinaDataset):
|
||||
"""
|
||||
Dataset class for the PINA dataset with :class:`~torch_geometric.data.Data`
|
||||
and :class:`~pina.graph.Graph` data.
|
||||
"""
|
||||
|
||||
def _create_graph_batch(self, data):
|
||||
class PinaGraphDataset(Dataset):
|
||||
def __init__(self, data_dict, automatic_batching=None):
|
||||
"""
|
||||
Create a LabelBatch object from a list of
|
||||
:class:`~torch_geometric.data.Data` objects.
|
||||
Initialize the instance by storing the conditions dictionary.
|
||||
|
||||
:param data: List of items to collate in a single batch.
|
||||
:type data: list[Data] | list[Graph]
|
||||
:return: LabelBatch object all the graph collated in a single batch
|
||||
disconnected graphs.
|
||||
:rtype: LabelBatch
|
||||
"""
|
||||
batch = LabelBatch.from_data_list(data)
|
||||
return batch
|
||||
|
||||
def create_batch(self, data):
|
||||
"""
|
||||
Create a Batch object from a list of :class:`~torch_geometric.data.Data`
|
||||
objects.
|
||||
|
||||
:param data: List of items to collate in a single batch.
|
||||
:type data: list[Data] | list[Graph]
|
||||
:return: Batch object.
|
||||
:rtype: :class:`~torch_geometric.data.Batch`
|
||||
| :class:`~pina.graph.LabelBatch`
|
||||
:param dict conditions_dict: A dictionary mapping condition names to
|
||||
their respective data. Each key represents a condition name, and the
|
||||
corresponding value is a dictionary containing the associated data.
|
||||
"""
|
||||
|
||||
if isinstance(data[0], Data):
|
||||
return self._create_graph_batch(data)
|
||||
return self._create_tensor_batch(data)
|
||||
# Store the conditions dictionary
|
||||
self.data = data_dict
|
||||
self.automatic_batching = (
|
||||
automatic_batching if automatic_batching is not None else True
|
||||
)
|
||||
|
||||
# Override _retrive_data method for graph handling
|
||||
def _retrive_data(self, data, idx_list):
|
||||
def __len__(self):
|
||||
return len(next(iter(self.data.values())))
|
||||
|
||||
def __getitem__(self, idx):
|
||||
"""
|
||||
Retrieve data from the dataset given a list of indices.
|
||||
Return the data at the given index in the dataset.
|
||||
|
||||
:param dict data: Dictionary containing the data.
|
||||
:param list[int] idx_list: List of indices to retrieve.
|
||||
:return: Dictionary containing the data at the given indices.
|
||||
:param int idx: Index.
|
||||
:return: A dictionary containing the data at the given index.
|
||||
:rtype: dict
|
||||
"""
|
||||
|
||||
if self.automatic_batching:
|
||||
# Return the data at the given index
|
||||
return {
|
||||
field_name: data[idx] for field_name, data in self.data.items()
|
||||
}
|
||||
return idx
|
||||
|
||||
def _getitem_from_list(self, idx_list):
|
||||
"""
|
||||
Return data from the dataset given a list of indices.
|
||||
|
||||
:param list[int] idx_list: List of indices.
|
||||
:return: A dictionary containing the data at the given indices.
|
||||
:rtype: dict
|
||||
"""
|
||||
|
||||
# Return the data from the current condition
|
||||
# If the data is a list of Data objects, create a Batch object
|
||||
# If the data is a list of torch.Tensor objects, create a torch.Tensor
|
||||
return {
|
||||
k: (
|
||||
self._create_graph_batch([v[i] for i in idx_list])
|
||||
if isinstance(v, list)
|
||||
else v[idx_list]
|
||||
)
|
||||
for k, v in data.items()
|
||||
field_name: [data[i] for i in idx_list]
|
||||
for field_name, data in self.data.items()
|
||||
}
|
||||
|
||||
@property
|
||||
def input(self):
|
||||
"""
|
||||
Return the input data for the dataset.
|
||||
|
||||
:return: Dictionary containing the input points.
|
||||
:rtype: dict
|
||||
"""
|
||||
return {k: v["input"] for k, v in self.conditions_dict.items()}
|
||||
|
||||
Reference in New Issue
Block a user