omniloader.data.datasets

Dataset adapters that yield raw sample dicts for OmniLoader.

The concrete datasets here (an HDF5-backed reader and an in-memory tensor dataset) implement the loose contract OmniLoader expects: __len__ returns the number of samples and __getitem__ returns a dict of named values. The DatasetSchema describing what a dataset provides is supplied alongside it to OmniLoader.

Classes

HDF5Dataset	Read samples from a group of an HDF5 file, one sub-group per sample.
DictTensorDataset	In-memory dataset over a dict of stacked tensors, keyed by sample index.

Module Contents

class omniloader.data.datasets.HDF5Dataset(h5_path: str | pathlib.Path, subset: str, keys: collections.abc.Sequence[str] | None = None, cache_size: int = 0, preload: bool = False)

Bases: torch.utils.data.Dataset

Read samples from a group of an HDF5 file, one sub-group per sample.

The file is expected to hold a group named subset (e.g. "train") whose children are per-sample groups, each containing the named value datasets (features, targets, masks, metadata). The file handle is opened lazily on first read so it can be per-process, but an HDF5 handle is not fork-safe: if anything reads a sample (or preload=True) before a DataLoader forks its workers, every worker inherits the same open handle and reads silently hang. Use num_workers > 0 only with omniloader.utils.seeding.seed_worker() passed as the DataLoader’s worker_init_fn (OmniLoader’s own loader/config helpers already do this), which clears the handle so each worker reopens its own.

Parameters:

• h5_path – Path to the .h5 file.

• subset – Name of the top-level group to read (e.g. "train").

• keys – Optional subset of value names to load per sample. When None, every dataset in the sample group is loaded.

• cache_size – If > 0, keep an in-process LRU cache of up to this many decoded samples to speed up repeated reads.

• preload – If True, decode every sample into memory up front (fastest for small datasets; uses the most memory).

Raises:

KeyError – If subset is not present in the file.

__len__() → int

Number of samples in the subset.

__getitem__(index: int) → dict[str, Any]

Load the sample at index as a dict of named values.

Parameters:

index – Position within the subset.

Returns:

A dict mapping each stored key to a tensor or string.

__getstate__() → dict[str, Any]

Drop the open file handle and cache when pickling (for workers).

reset_for_worker() → None

Drop any inherited file handle and cache so this process reopens its own.

Called by omniloader.utils.seeding.seed_worker() (the DataLoader’s worker_init_fn) inside each worker, since a fork-started worker inherits the parent’s already-open (and not fork-safe) handle directly without going through __getstate__().

class omniloader.data.datasets.DictTensorDataset(tensors: collections.abc.Mapping[str, torch.Tensor], metadata: collections.abc.Mapping[str, collections.abc.Sequence[Any]] | None = None)

Bases: torch.utils.data.Dataset

In-memory dataset over a dict of stacked tensors, keyed by sample index.

Useful for tests and small experiments: each entry maps a name to a tensor whose leading dimension is the sample axis. __getitem__ returns the per-sample slice of every tensor as a dict.

Parameters:

• tensors – Mapping of value name to a tensor of shape (N, ...) where N is the number of samples (equal across all tensors).

• metadata – Optional mapping of metadata name to a per-sample sequence (e.g. dataset name or sample key) of length N.

Raises:

ValueError – If the tensors disagree on the number of samples.

__len__() → int

Number of samples.

__getitem__(index: int) → dict[str, Any]

Return the per-sample slice of every tensor plus metadata.

Parameters:

index – Sample position.

Returns:

A dict of the sample’s values and metadata.