"""
Base dataset classes for use with BaseArchitecture
"""
from __future__ import annotations
import logging as log
from types import ModuleType
from itertools import repeat
from typing import (
Any,
Self,
Generic,
Literal,
Sequence,
Iterator,
Protocol,
TypeVar,
cast,
overload,
)
import torch
import numpy as np
from torch import Tensor
from torch.utils.data import Dataset, DataLoader, Subset
from numpy import ndarray
from netloader.utils.types import (
DataLike,
ArrayLike,
DataT,
ArrayT,
DatasetT,
DataListT,
ArrayCT,
P,
)
InArrayT_contra = TypeVar('InArrayT_contra', bound=DataLike, contravariant=True)
OutArrayT_co = TypeVar('OutArrayT_co', bound=DataLike, covariant=True)
[docs]
class ApplyFunc(Protocol[InArrayT_contra, P, OutArrayT_co]):
"""
Protocol for functions that can be applied to Data objects in the apply method of Data and
DataList.
"""
[docs]
def __call__(self, data: InArrayT_contra, /, *args: P.args, **kwargs: P.kwargs) -> OutArrayT_co:
"""
Parameters
----------
data : DataT
Data to apply the function to
*args
Optional arguments to pass to the function
**kwargs
Optional keyword arguments to pass to the function
Returns
-------
DataT
Result of applying the function to the data
"""
[docs]
class Data(Generic[ArrayCT]):
"""
Stores data with uncertainties for uncertainty handling in BaseArchitecture.
Attributes
----------
data : ArrayCT
Data
uncertainty : ArrayCT | None
Data uncertainty
"""
[docs]
def __init__(self, data: ArrayCT, uncertainty: ArrayCT | None = None) -> None:
"""
Parameters
----------
data : ArrayCT
Data
uncertainty : ArrayCT | None, Optional
Data uncertainty
"""
self.shape: tuple[int, ...] = tuple(data.shape)
self.data: ArrayCT = data
self.uncertainty: ArrayCT | None = uncertainty
[docs]
def __len__(self) -> int:
"""
Returns the number of samples in the data.
Returns
-------
int
Number of samples in the data
"""
return len(self.data)
[docs]
def __getitem__(self, idx: int | slice) -> Data[ArrayCT]:
"""
Gets a subset of the data and uncertainty.
Parameters
----------
idx : int | slice
Index or slice to get
Returns
-------
Data[ArrayCT]
Data with subset of data and uncertainty
"""
return Data(
self.data[idx],
uncertainty=self.uncertainty[idx] if self.uncertainty is not None else None
)
def __repr__(self) -> str:
return (f'Data(shape={self.shape}, type={self.data.dtype}, '
f'uncertainty={self.uncertainty is not None})')
@staticmethod
def _apply(
data: ArrayCT,
func: str | ApplyFunc[ArrayCT, ..., ArrayCT],
*args: Any,
types: tuple[type[DataLike], ...] = (Tensor, ndarray),
**kwargs: Any) -> ArrayCT:
"""
Applies a function to the data.
Parameters
----------
data : ArrayCT
Data to apply the function to
func : str | ApplyFunc[ArrayCT, ..., ArrayCT]
Function, or method if string, to apply to the data and uncertainty
*args
Arguments to pass to the function or method
types : tuple[type[ArrayLike], ...]
Types to apply the function or method to, if data and uncertainty is not an instance of
types, then data and uncertainty are returned unchanged, default = (Tensor, ndarray)
**kwargs
Keyword arguments to pass to the function or method
Returns
-------
ArrayCT
Result of applying the function to the data
"""
if data not in types:
return data
if isinstance(func, str):
return getattr(data, func)(*args, **kwargs)
return func(data, *args, **kwargs)
@staticmethod
@overload
def collate(data: list[Data[ArrayCT]], *, data_field: Literal[True]) -> Data[ArrayCT]: ...
@staticmethod
@overload
def collate(data: list[Data[ArrayCT]], *, data_field: Literal[False]) -> ArrayCT: ...
@staticmethod
@overload
def collate(data: list[Data[ArrayCT]], *, data_field: bool) -> ArrayCT | Data[ArrayCT]: ...
[docs]
@staticmethod
def collate(
data: list[Data[ArrayCT]],
*,
data_field: bool = True) -> ArrayCT | Data[ArrayCT]:
"""
Collates a list of Data objects into a single Data object.
Parameters
----------
data : list[Data[ArrayCT]]
List of Data objects to collate
data_field : bool, Optional
If True, returns a Data object, else returns the collated data as an ArrayLike,
default = True
Returns
-------
ArrayCT | Data[ArrayCT]
Collated ArrayLike or Data object
"""
module: ModuleType = torch if isinstance(data[0].data, Tensor) else np
new_data: ArrayCT
datum: Data[Any]
if data[0].uncertainty is None:
new_data = module.concat([datum.data for datum in data])
return Data(new_data) if data_field else new_data
new_data = module.concat([datum.concat() for datum in data])
return Data(*new_data.swapaxes(0, 1)) if data_field else new_data
[docs]
def apply(
self,
func: str | ApplyFunc[ArrayCT, ..., ArrayCT],
*args: Any,
types: tuple[type[DataLike], ...] = (Tensor, ndarray),
**kwargs: Any) -> Self:
"""
Applies a function or method to the data and uncertainty.
Parameters
----------
func : str | ApplyFunc[ArrayCT]
Function, or method if string, to apply to the data and uncertainty
*args
Arguments to pass to the function or method
types : tuple[type[ArrayLike], ...]
Types to apply the function or method to, if data and uncertainty is not an instance of
types, then data and uncertainty are returned unchanged, default = (Tensor, ndarray)
**kwargs
Keyword arguments to pass to the function or method
Returns
-------
Self
Self with function or method applied to the data and uncertainty
"""
if not isinstance(self.data, types):
return self
self.data = self._apply(self.data, func, *args, types=types, **kwargs)
self.uncertainty = None if self.uncertainty is None else self._apply(
self.uncertainty,
func,
*args,
types=types,
**kwargs,
)
return self
[docs]
def clone(self) -> Data[ArrayCT]:
"""
Clones the data and uncertainty.
Identical to Data.copy().
Returns
-------
Data[ArrayCT]
Cloned Data
"""
data: ArrayCT
uncertainty: ArrayCT | None
if isinstance(self.data, ndarray):
data = self.data.copy()
uncertainty = self.uncertainty.copy() if self.uncertainty is not None else None
else:
data = self.data.clone()
uncertainty = self.uncertainty.clone() if self.uncertainty is not None else None
return Data(data, uncertainty=uncertainty)
[docs]
def concat(self, dim: int = 0) -> ArrayCT:
"""
Concatenates data and uncertainty for passing into a network.
Parameters
----------
dim : int, Optional
Dimension to concatenate along, default = 0
Returns
-------
ArrayCT
Data and uncertainty concatenated along the specified dimension if uncertainty is not
None, else just data
"""
module: ModuleType = torch if isinstance(self.data, Tensor) else np
kwargs: dict[str, int] = {'dim' if isinstance(self.data, Tensor) else 'axis': dim}
if self.uncertainty is not None:
return module.concat((self.data, self.uncertainty), **kwargs)
return self.data
copy = clone
[docs]
def cpu(self) -> Self:
"""
Moves data and uncertainty to CPU if they are Tensors.
Returns
-------
Self
Self with data and uncertainty on CPU
"""
if isinstance(self.data, Tensor):
self.data = self.data.cpu()
if isinstance(self.uncertainty, Tensor):
self.uncertainty = self.uncertainty.cpu()
return self.apply('cpu')
[docs]
def detach(self) -> Self:
"""
Detaches data and uncertainty from the computation graph if they are Tensors.
Returns
-------
Self
Self with data and uncertainty detached from the computation graph
"""
if isinstance(self.data, Tensor):
self.data = self.data.detach()
if isinstance(self.uncertainty, Tensor):
self.uncertainty = self.uncertainty.detach()
return self
[docs]
def numpy(self) -> Data[ndarray]:
"""
Converts data and uncertainty to numpy arrays if they are Tensors.
Returns
-------
Data[ndarray]
New Data object with data and uncertainty as numpy arrays
"""
return Data(
self.data.cpu().numpy() if isinstance(self.data, Tensor) else self.data,
uncertainty=self.uncertainty.cpu().numpy() if isinstance(self.uncertainty, Tensor) else
self.uncertainty
)
[docs]
def tensor(self) -> Data[Tensor]:
"""
Converts data and uncertainty to tensors if they are numpy arrays.
Returns
-------
Data[Tensor]
New Data object with data and uncertainty as tensors
"""
return Data(
torch.from_numpy(self.data) if isinstance(self.data, ndarray) else self.data,
uncertainty=torch.from_numpy(self.uncertainty) if isinstance(self.uncertainty, ndarray)
else self.uncertainty
)
[docs]
def to(self, *args: Any, **kwargs: Any) -> Self:
"""
Move and/or cast the parameters and buffers.
Parameters
----------
*args
Arguments to pass to the to method of data and uncertainty
**kwargs
Keyword arguments to pass to the to method of data and uncertainty
Returns
-------
Self
Self with data and uncertainty moved and/or cast
"""
if isinstance(self.data, Tensor):
self.data = self.data.to(*args, **kwargs)
if isinstance(self.uncertainty, Tensor):
self.uncertainty = self.uncertainty.to(*args, **kwargs)
return self
[docs]
class DataList(Generic[DataT]):
"""
A list that stores tensors, arrays, or Datas and provides batch operations.
"""
[docs]
def __init__(self, data: list[DataT]) -> None:
"""
Parameters
----------
data : list[DataT]
List of Data, Tensor, or ndarray objects
"""
self._data: list[DataT] = data
if any(len(data[0]) != len(datum) for datum in data[1:]):
raise ValueError(f'All elements in DataList must have the same length, got lengths: '
f'{[len(datum) for datum in data]}')
[docs]
def __add__(self, other: DataT | DataList[DataT]) -> DataList[DataT]:
"""
Adds another DataList or data to the DataList.
Parameters
----------
other : DataT | DataList[DataT]
DataList or single element to add to the DataList
Returns
-------
DataList[DataT]
New DataList with the other DataList or data added to the end of the DataList
"""
if len(self) != len(other):
raise ValueError(f'Lengths of DataList ({len(self)}) and {type(other).__name__} '
f'({len(other)}) must be the same')
return DataList(self._data + (list(other) if isinstance(other, DataList) else [other]))
[docs]
def __getitem__(self, idx: int | slice) -> DataList[DataT]:
"""
Gets a subset of each element in the DataList.
Parameters
----------
idx: int | slice
Index or slice to get
Returns
-------
DataList[DataT]
DataList with subset of each element in the DataList
"""
return self.get(idx, list_=False)
[docs]
def __iter__(self) -> Iterator[DataT]:
"""
Iterates over each element in the DataList.
Returns
-------
Iterator[DataT]
Iterator over each element in the DataList
"""
return self.iter(list_=True)
[docs]
def __len__(self) -> int:
"""
Returns the number of samples in each element.
Returns
-------
int
Number of samples in each element
"""
return self.len(list_=False)
def __repr__(self) -> str:
"""
Returns a string representation of the DataList.
Returns
-------
str
String representation of the DataList
"""
return (f'DataList(shapes={[tuple(datum.shape) for datum in self._data]}, '
f'types={[datum.__class__.__name__ for datum in self._data]})')
@staticmethod
def _apply(
data: DataT,
func: str | ApplyFunc[DataT, ..., DataT],
*args: Any,
types: tuple[type[DataLike], ...] = (ndarray, Tensor, Data),
**kwargs: Any) -> DataT:
"""
Applies a function or method to the data if it is an instance of types, else returns the
data unchanged.
Parameters
----------
data : DataT
Data to apply function to
func : str | ApplyFunc[ArrayT]
Function, or method if string, to apply to the data
*args
Arguments to pass to the function or method
types : tuple[type[DataLike], ...]
Types to apply the function or method to, if data is not an instance of types, then
data is returned unchanged, default = (ndarray, Tensor, Data)
**kwargs
Keyword arguments to pass to the function or method
Returns
-------
DataT
Result of applying the function or method to the data if it is an instance of types,
else returns the data unchanged
"""
if not isinstance(data, types):
return data
if isinstance(func, str):
return getattr(data, func)(*args, **kwargs)
return func(data, *args, **kwargs)
@staticmethod
@overload
def collate(
data: list[DataList[DataT]],
*,
data_field: Literal[True]) -> DataList[DataT]: ...
@staticmethod
@overload
def collate(
data: list[DataList[DataT]],
*,
data_field: Literal[False]) -> DataList[ArrayT]: ...
@staticmethod
@overload
def collate(
data: list[DataList[DataT]],
*,
data_field: bool) -> DataList[ArrayT] | DataList[DataT]: ...
[docs]
@staticmethod
def collate(
data: list[DataList[DataT]],
*,
data_field: bool = True) -> DataList[ArrayT] | DataList[DataT]:
"""
Collates a list of DataList objects into a single DataList object.
Parameters
----------
data : list[DataList[DataT]]
List of DataList objects to collate
data_field : bool, Optional
If True, collates Data elements into Data object, else collates into ArrayLike,
default = True
Returns
-------
DataList[ArrayT] | DataList[DataT]
Collated DataList object
"""
i: int
element_data: list[DataLike]
new_data: list[DataLike] = []
datum: DataList[DataT]
for i in range(data[0].len(True)):
element_data = [datum.get(i, list_=True) for datum in data]
new_data.append(
Data.collate(cast(list[Data], element_data), data_field=data_field)
if isinstance(element_data[0], Data) else
torch.concat(cast(list[Tensor], element_data))
if isinstance(element_data[0], Tensor) else
np.concat(cast(list[ndarray], element_data)),
)
return cast(DataList[ArrayT] | DataList[DataT], DataList(new_data))
[docs]
def append(self, data: DataT) -> None:
"""
Appends an element to the DataList.
Parameters
----------
data : DataT
Element to append to the DataList
"""
if self.len(list_=False) != len(data):
raise ValueError(f'Element to append must have the same length as the DataList, got '
f'lengths: {self.len(list_=False)} and {len(data)}')
self._data.append(data)
[docs]
def apply(
self,
func: str | Sequence[str | ApplyFunc[DataT, ..., DataT]] | ApplyFunc[DataT, ..., DataT],
*args: Any,
types: tuple[type[DataLike], ...] = (Data, Tensor, ndarray),
**kwargs: Any) -> Self:
"""
Applies a function or method to each element in the DataList.
Parameters
----------
func : str | Sequence[str | ApplyFunc[ArrayT]] | ApplyFunc[DataT]
Function(s), or method(s) if string, to apply to each element in the DataList
*args
Arguments to pass to the function or method
types : tuple[type[DataLike], ...]
Types to apply the function or method to, if an element is not an instance of types,
then that element is returned unchanged, default = (Data, Tensor, ndarray)
**kwargs
Keyword arguments to pass to the function or method
Returns
-------
Self
Self with function or method applied to each element in the DataList
"""
datum: DataT
func_: str | ApplyFunc[DataT, ..., DataT]
if isinstance(func, list) and len(func) != len(self._data):
raise ValueError(f'Length of function list ({len(func)}) must match length of DataList '
f'({len(self._data)})')
self._data = [
self._apply(datum, func_, *args, types=types, **kwargs)
for datum, func_ in zip(self._data, func if isinstance(func, list) else repeat(func))
]
return self
[docs]
def clone(self) -> DataList[DataT]:
"""
Clones the DataList.
Identical to DataList.copy().
Returns
-------
DataList[DataT]
Cloned DataList
"""
data: DataT
return DataList([
cast(DataT, data.clone() if isinstance(data, Tensor) else data.copy())
for data in self._data
])
copy = clone
[docs]
def cpu(self) -> Self:
"""
Moves all tensors to CPU
Returns
-------
Self
Self with all tensors moved to CPU
"""
return self.apply('cpu', types=(Tensor, Data))
[docs]
def detach(self) -> Self:
"""
Detaches all tensors from the computation graph
Returns
-------
Self
Self with all tensors detached from the computation graph
"""
return self.apply('detach', types=(Tensor, Data))
[docs]
def extend(self, data: Sequence[DataT] | DataList[DataT]) -> None:
"""
Extends the DataList by appending elements from the iterable.
Parameters
----------
data : Sequence[DataT] | DataList[DataT]
Elements to extend the DataList with
"""
datum: DataT
for datum in data:
self.append(datum)
@overload
def get(self, idx: int, list_: Literal[True]) -> DataT: ...
@overload
def get(self, idx: slice, list_: Literal[True]) -> list[DataT]: ...
@overload
def get(self, idx: int | slice, list_: Literal[False]) -> DataList[DataT]: ...
@overload
def get(self, idx: int, list_: bool) -> DataT | DataList[DataT]: ...
@overload
def get(self, idx: slice, list_: bool) -> list[DataT] | DataList[DataT]: ...
[docs]
def get(self, idx: int | slice, list_: bool = False) -> DataT | list[DataT] | DataList[DataT]:
"""
Gets a subset of the DataList or a subset of each element in the DataList.
Parameters
----------
idx : int | slice
Index or slice to get
list_ : bool, Optional
If True, returns a subset of the DataList, else returns a subset of each element in the
DataList, default = False
Returns
-------
DataT | list[DataT] | DataList[DataT]
Subset of the DataList or subset of each element in the DataList
"""
data: DataT
if list_:
return self._data[idx]
return DataList(cast(list[DataT], [data[idx] for data in self]))
[docs]
def get_data(self) -> list[DataT]:
"""
Gets the underlying data list.
Returns
-------
list[DataT]
Underlying data list
"""
return self._data
[docs]
def insert(self, idx: int, data: DataT) -> None:
"""
Inserts an element into the DataList at the specified index.
Parameters
----------
idx : int
Index to insert the element at
data : DataT
Element to insert into the DataList
"""
if self.len(list_=False) != len(data):
raise ValueError(f'Element to insert must have the same length as the DataList, got '
f'lengths: {self.len(list_=False)} and {len(data)}')
self._data.insert(idx, data)
@overload
def iter(self, list_: Literal[True]) -> Iterator[DataT]: ...
@overload
def iter(self, list_: Literal[False]) -> Iterator[DataList[DataT]]: ...
@overload
def iter(self, list_: bool) -> Iterator[DataT] | Iterator[DataList[DataT]]: ...
[docs]
def iter(self, list_: bool = True) -> Iterator[DataT] | Iterator[DataList[DataT]]:
"""
Iterates over each element in the DataList.
Parameters
----------
list_ : bool, Optional
If True, iterates over the DataList, else iterates over each element in the DataList,
default = True
Returns
-------
Iterator[DataT] | Iterator[DataList[DataT]]
Iterator over each element in the DataList
"""
i: int
for i in range(self.len(True)):
yield self.get(i, list_=list_)
[docs]
def len(self, list_: bool = True) -> int:
"""
Gets the length of the DataList or the length of each element in the DataList.
Parameters
----------
list_ : bool, Optional
If True, returns the length of the DataList, else returns the length of each element
in the DataList, default = True
Returns
-------
int
Length of the DataList or length of each element in the DataList
"""
return len(self._data) if list_ else len(self._data[0])
[docs]
def numpy(self) -> DataList[ndarray | Data[ndarray]]:
"""
Converts all tensors to numpy arrays.
Returns
-------
DataList[ndarray | Data[ndarray]]
DataList with all tensors converted to numpy arrays
"""
data: DataT
return DataList(cast(list[ndarray | Data[ndarray]], [
data.cpu().numpy() if hasattr(data, 'cpu') else data for data in self
]))
[docs]
def tensor(self) -> DataList[Tensor | Data[Tensor]]:
"""
Converts all numpy arrays to tensors.
Returns
-------
DataList[Tensor | Data[Tensor]]
DataList with all numpy arrays converted to tensors
"""
data: DataT
return DataList(cast(list[Tensor | Data[Tensor]], [
torch.from_numpy(data) if isinstance(data, ndarray) else
data.tensor() if isinstance(data, Data) else data for data in self
]))
[docs]
def to(self, *args: Any, **kwargs: Any) -> Self:
"""
Move and/or cast the parameters and buffers.
Parameters
----------
*args
Arguments to pass to the to method of each element in the DataList
**kwargs
Keyword arguments to pass to the to method of each element in the DataList
Returns
-------
Self
Self with all elements moved and/or cast
"""
return self.apply('to', *args, types=(Tensor, Data), **kwargs)
class BaseDatasetMeta(type):
"""
Automatically creates an index for each sample in the dataset after the dataset has been
initialised.
"""
def __call__(cls: type[DatasetT], *args: Any, **kwargs: Any) -> DatasetT:
"""
Parameters
----------
cls : type[DatasetT]
Class that inherited BaseDatasetMeta
*args
Optional arguments to pass to BaseDataset class
**kwargs
Optional keyword arguments to pass to BaseDataset class
Returns
-------
DatasetT
Dataset instance from the class that inherited BaseDatasetMeta
"""
instance: DatasetT = type.__call__(cls, *args, **kwargs)
if hasattr(instance, 'idxs') and instance.idxs.dtype != np.int_:
raise ValueError(f'idxs attribute already exists and does not have type int '
f'({instance.idxs.dtype}), idxs attribute must be reserved for sample '
f'index')
if not hasattr(instance, 'high_dim') or instance.high_dim is None:
raise ValueError(f'{instance.__class__.__name__} has no high_dim attribute which is '
f'required by BaseDatasetMeta for creating idxs attribute')
if len(instance.idxs) == 0:
instance.idxs = np.arange(len(instance.high_dim))
elif len(instance.idxs) != len(instance.high_dim):
log.getLogger(__name__).warning(f'Length of idxs ({len(instance.idxs)}) and length of '
f'high_dim ({len(instance.high_dim)}) does not match, '
f'idxs will be sent to a range of high_dim length')
instance.idxs = np.arange(len(instance.high_dim))
for attribute in ('extra', 'low_dim', 'high_dim'):
if (getattr(instance, attribute) is not None and
len(getattr(instance, attribute)) != len(instance.idxs)):
raise ValueError(f'Length of attribute {attribute} '
f'({len(getattr(instance, attribute))}) and idxs '
f'({len(instance.idxs)}) does not match')
return instance
[docs]
class BaseDataset(
Dataset[tuple[int, DataListT, DataListT, Any]],
Generic[DataListT],
metaclass=BaseDatasetMeta):
"""
Base dataset class for use with BaseNetwork.
Attributes
----------
extra : list[Any] | ArrayLike | None
Additional data for each sample in the dataset of length N with shape (N,...) and type Any
idxs : ndarray
Index for each sample in the dataset with shape (N) and type int
low_dim : DataListT | None
Low dimensional data for each sample in the dataset with shape (N,...)
high_dim : DataListT | None
High dimensional data for each sample in the dataset with shape (N,...), this is required
"""
def __init__(self) -> None:
super().__init__()
self.extra: list[Any] | ArrayLike | None = None
self.idxs: ndarray = np.array([], dtype=np.int_)
self.low_dim: DataListT | None = None
self.high_dim: DataListT | None = None
[docs]
def __len__(self) -> int:
"""
Returns the number of samples in the dataset
Returns
-------
int
Number of samples in the dataset
"""
return len(self.idxs)
[docs]
def __getitem__(self, idx: int) -> tuple[int, DataListT, DataListT, Any]:
"""
Parameters
----------
idx : int
Sample index
Returns
-------
tuple[int, DataListT, DataListT, Any]
Sample index, low dimensional data, high dimensional data, and extra data
"""
return self.idxs[idx], self.get_low_dim(idx), self.get_high_dim(idx), self.get_extra(idx)
[docs]
def get_high_dim(self, idx: int) -> DataListT:
"""
Gets a high dimensional sample of the given index
Parameters
----------
idx : int
Sample index
Returns
-------
DataListT
High dimensional sample
"""
assert self.high_dim is not None
return cast(DataListT, self.high_dim[idx])
[docs]
def get_low_dim(self, idx: int) -> DataListT:
"""
Gets a low dimensional sample of the given index
Parameters
----------
idx : int
Sample index
Returns
-------
DataListT
Low dimensional sample
"""
return cast(DataListT, torch.tensor(()) if self.low_dim is None else self.low_dim[idx])
[docs]
def step(self, epoch: float) -> None:
"""
Step the dataset for each training iteration.
Parameters
----------
epoch : float
Current epoch number
"""
@overload
def data_collation(
data: list[ArrayCT] | ArrayCT,
*,
data_field: bool) -> ArrayCT: ...
@overload
def data_collation(
data: DataList[DataT],
*,
data_field: bool) -> DataList[DataT]: ...
@overload
def data_collation(
data: list[DataList[DataT]],
*,
data_field: Literal[True]) -> DataList[DataT]: ...
@overload
def data_collation(
data: list[DataList[DataT]],
*,
data_field: Literal[False]) -> DataList[ArrayT]: ...
@overload
def data_collation(
data: list[Data[ArrayCT]],
*,
data_field: Literal[True]) -> Data[ArrayCT]: ...
@overload
def data_collation(
data: list[Data[ArrayCT]],
*,
data_field: Literal[False]) -> ArrayCT: ...
[docs]
def data_collation( # type: ignore
data: list[ArrayCT] | list[Data[ArrayCT]] | list[DataList[DataT]] | ArrayCT |
DataList[DataT],
*,
data_field: bool = True,
) -> ArrayCT | Data[ArrayCT] | DataList[ArrayT] | DataList[DataT]:
"""
Collates a list of ArrayLike, Data, or DataList objects into a single object.
Parameters
----------
data : list[ArrayCT] | list[Data[ArrayCT]] | list[DataList[DataT]] | ArrayCT | DataList[DataT]
List of ArrayLike, Data, or DataList objects to collate, or if ArrayLike or DataList,
will return as is, with shape (N,...)
data_field : bool, Optional
If Datas should return Data else ArrayLike, default = True
Returns
-------
ArrayCT | Data[ArrayCT] | DataList[ArrayT] | DataList[DataT]
Collated Data or DataList object, or ArrayLike if data_field is False or input is
ArrayLike, with shape (N,...)
"""
if isinstance(data, (Tensor, ndarray, DataList)):
return data
if isinstance(data[0], (Tensor, ndarray)):
return torch.concat(cast(list[Tensor], data)) if isinstance(data[0], Tensor) else \
np.concat(cast(list[ndarray], data))
# data = cast(list[Data] | list[DataList], data)
return Data.collate(cast(list[Data], data), data_field=data_field) \
if isinstance(data[0], Data) else \
DataList.collate(cast(list[DataList], data), data_field=data_field)
@overload
def loader_init(
dataset: DatasetT,
*,
return_idxs: Literal[False],
batch_size: int = ...,
ratios: list[float] | tuple[float, ...] | None = ...,
idxs: list[ndarray] | tuple[ndarray, ...] | ndarray | None = ...,
**kwargs: Any) -> tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...]: ...
@overload
def loader_init(
dataset: DatasetT,
*,
return_idxs: Literal[True],
batch_size: int = ...,
ratios: list[float] | tuple[float, ...] | None = ...,
idxs: list[ndarray] | tuple[ndarray, ...] | ndarray | None = ...,
**kwargs: Any) -> tuple[
tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...],
tuple[Sequence[int], ...]]: ...
[docs]
def loader_init(
dataset: DatasetT,
*,
return_idxs: bool = False,
batch_size: int = 64,
ratios: list[float] | tuple[float, ...] | None = None,
idxs: list[ndarray] | tuple[ndarray, ...] | ndarray | None = None,
**kwargs: Any,
) -> (tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...] |
tuple[
tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...],
tuple[Sequence[int], ...]]):
# pylint: disable=line-too-long
"""
Initialises data loaders from a subset of the dataset with the given ratios.
Parameters
----------
dataset : DatasetT
Dataset to create data loaders from
return_idxs : bool, Optional
If the indexes for each data loader should be returned, default = False
batch_size : int, Optional
Batch size when sampling from the data loaders, default = 64
ratios : list[float] | tuple[float, ...] | None, Optional
Ratios of length M to split up the dataset into subsets, if idxs is provided, dataset will
first be split up using idxs and ratios will be used on the remaining samples,
default = (0.8,0.2)
idxs : list[ndarray] | tuple[ndarray, ...] | ndarray | None, Optional
Dataset indexes for creating the subsets with shape (N,S), where N is the number of subsets
and S is the number of samples in each subset
**kwargs : Any
Optional keyword arguments to pass to DataLoader
Returns
-------
tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...] | tuple[tuple[DataLoader[tuple[int, DataListT, DataListT, Any]], ...], tuple[Sequence[int], ...]]
Data loaders for each subset of length N + M and optionally indexes used for each subset
"""
# pylint: enable=line-too-long
num: int
slice_: float | ndarray
loaders: list[DataLoader[tuple[int, DataListT, DataListT, Any]]] = []
data_idxs: ndarray = np.arange(len(dataset))
loader: DataLoader[tuple[int, DataListT, DataListT, Any]]
idxs = [] if idxs is None else list(idxs) \
if isinstance(idxs, (tuple, list)) or np.ndim(idxs) > 1 else [idxs]
ratios = (0.8, 1) if ratios is None else list(np.cumsum(np.array(ratios) / np.sum(ratios)))
np.random.shuffle(data_idxs)
for slice_ in idxs + list(ratios):
if isinstance(slice_, (int, float)):
num = max(int(len(data_idxs) * slice_), 1)
slice_ = data_idxs[:num]
if not np.isin(data_idxs, slice_).any():
continue
loaders.append(DataLoader(
Subset(
cast(Dataset[tuple[int, DataListT, DataListT, Any]], dataset),
data_idxs[np.isin(data_idxs, slice_)].tolist(),
),
batch_size=batch_size,
**{'shuffle': True} | kwargs,
))
data_idxs = np.delete(data_idxs, np.isin(data_idxs, slice_))
if return_idxs:
return tuple(loaders), tuple(cast(Subset, loader.dataset).indices for loader in loaders)
return tuple(loaders)
__all__ = ['ApplyFunc', 'Data', 'DataList', 'BaseDataset', 'data_collation', 'loader_init']