omniloader.config¶

Configuration and reproducibility helpers.

OmniConfig bundles the knobs a training run needs — seed, batch size, worker count, the mixing strategy and the normalization/augmentation pipeline — and loads from a dict, a JSON file or a YAML file so experiments stay declarative. seed_everything() seeds Python, NumPy and PyTorch in one call.

Attributes¶

`STRATEGIES`
`COLLATES`

Classes¶

OmniConfig

Declarative configuration for an OmniLoader training run.

Functions¶

seed_everything(→ int)

Seed Python, NumPy and PyTorch RNGs for reproducibility.

Module Contents¶

omniloader.config.STRATEGIES: dict[str, type[omniloader.sampling.strategies.MixingStrategy]]¶

omniloader.config.COLLATES: tuple[str, ...] = ('unified', 'dynamic', 'mixup')¶

omniloader.config.seed_everything(seed: int) → int[source]¶

Seed Python, NumPy and PyTorch RNGs for reproducibility.

Parameters:: seed – The seed value.
Returns:: The seed, for convenience.

class omniloader.config.OmniConfig[source]¶

Declarative configuration for an OmniLoader training run.

Parameters:

seed – Global RNG seed.
batch_size – Batch size for the DataLoader.
num_workers – Number of worker processes for the DataLoader.
pad_features – Whether the unifier pads/crops sequence values.
strategy – Name of the mixing strategy (a key of STRATEGIES).
strategy_kwargs – Extra keyword arguments forwarded to the strategy (e.g. temperature, weights, samples_per_dataset).
transforms – Ordered list of transform config dicts (each with a name key) applied during training. Stats-based normalizers may reference a saved stats file via stats_path (or union_stats_path) instead of inlining stats. See build_transform().
eval_transforms – Transforms applied during evaluation. When None, only the non-augmenting (train-agnostic) transforms in transforms, such as normalize, take effect.
datasets – Declarative dataset entries (adapter + args + schema) built by build_datasets(). See omniloader.data.factory.
subsample – Optional per-dataset subsampling entries (one per dataset, or None to skip a dataset), each a dict for SubsampleConfig. Built by build_subsample() and fed to build_strategy().
collate – Collate function name — one of COLLATES ("unified", "dynamic" or "mixup"). See build_collate(). "dynamic" only reduces padding with pad_features set to False (under a fixed time_dim it behaves like "unified").
collate_kwargs – Extra keyword arguments for the collator. For "dynamic" an optional keys list; for "mixup" a base collate name plus alpha/mode/p.
bucketing – Optional length-bucketing entry with a key (sequence name to bucket by) and optional batch_size/bucket_multiplier/ drop_last/shuffle. When set, build_dataloader() batches with LengthBucketBatchSampler. Only meaningful with pad_features=False (a no-op under a fixed time_dim).
num_replicas – Distributed replica count for the sampler. None infers it from torch.distributed.
rank – This process’s distributed rank. None infers it.
drop_last – Drop the uneven tail when sharding across replicas.
pin_memory – Copy tensors into pinned memory before returning them (faster host-to-GPU transfer). Forwarded to the DataLoader.
persistent_workers – Keep worker processes alive between epochs. Only takes effect when num_workers > 0.
prefetch_factor – Batches prefetched per worker. Only applied when num_workers > 0; None uses the PyTorch default.

classmethod from_dict(data: dict[str, Any]) → OmniConfig[source]¶

Build a config from a plain dict, ignoring unknown keys.

Parameters:: data – Mapping of field names to values.
Returns:: The populated config.

classmethod from_yaml(path: str | pathlib.Path) → OmniConfig[source]¶: Load a config from a YAML file.

classmethod from_json(path: str | pathlib.Path) → OmniConfig[source]¶: Load a config from a JSON file.

classmethod from_file(path: str | pathlib.Path) → OmniConfig[source]¶

Load a config from a .json, .yaml or .yml file by extension.

Parameters:: path – Path to a config file.
Returns:: The populated config.
Raises:: ValueError – If the file extension is not recognised.

to_dict() → dict[str, Any][source]¶: Return the config as a plain dict.

build_strategy(sizes: collections.abc.Sequence[int], subsample: collections.abc.Sequence[omniloader.sampling.strategies.SubsampleConfig | None] | None = None) → omniloader.sampling.strategies.MixingStrategy[source]¶

Instantiate the configured mixing strategy for the given sizes.

Parameters:

sizes – Per-dataset sample counts.
subsample – Optional per-dataset subsample configs.

Returns:

The configured MixingStrategy.

Raises:

ValueError – If strategy is not a known strategy name.

build_transform(schema: omniloader.schema.spec.UnifiedSchema, training: bool = True) → omniloader.transforms.Compose | None[source]¶

Build the transform pipeline for a split.

Parameters:

schema – The unified schema, injected into schema-aware transforms.
training – When True, uses transforms; otherwise uses eval_transforms if set, else transforms (whose augmentations self-skip in eval mode).

Returns:

A Compose, or None if empty.

build_datasets() → tuple[list[omniloader.loader.SizedDataset], list[omniloader.schema.spec.DatasetSchema]][source]¶

Construct the datasets/schemas declared in datasets.

Returns:: A (datasets, schemas) tuple. See omniloader.data.factory.build_datasets().

build_subsample() → list[omniloader.sampling.strategies.SubsampleConfig | None] | None[source]¶

Build the per-dataset SubsampleConfig list from subsample.

Returns:: One SubsampleConfig (or None) per dataset, or None when subsample is unset.

build_collate(schema: omniloader.schema.spec.UnifiedSchema, training: bool = True) → collections.abc.Callable[[collections.abc.Sequence[dict[str, Any]]], dict[str, Any]][source]¶

Build the collate function named by collate.

Parameters:

schema – The unified schema, injected into schema-aware collators.
training – When False, the batch-level "mixup" augmentation is disabled and its base collator is used instead.

Returns:

A callable turning a list of samples into a batched dict.

Raises:

ValueError – If collate (or a mixup base) is unknown.

build_loader(datasets: collections.abc.Sequence[omniloader.loader.SizedDataset] | None = None, schemas: collections.abc.Sequence[omniloader.schema.spec.DatasetSchema] | None = None, training: bool = True) → omniloader.loader.OmniLoader[source]¶

Build a fully configured OmniLoader.

Parameters:

datasets – Source datasets. When None, they are built from datasets via build_datasets().
schemas – Per-dataset schemas. Required when datasets is given; ignored (rebuilt) when datasets is None.
training – Whether augmentations are active (False for val/test).

Returns:

The loader with its transform pipeline wired in.

Raises:

ValueError – If datasets is given without schemas.

build_dataloader(datasets: collections.abc.Sequence[omniloader.loader.SizedDataset] | None = None, schemas: collections.abc.Sequence[omniloader.schema.spec.DatasetSchema] | None = None, training: bool = True, epoch: int = 0) → torch.utils.data.DataLoader[source]¶

Assemble a ready-to-iterate DataLoader entirely from this config.

Wires together the loader, mixing strategy, distributed sampler, collate function and (optionally) length bucketing. Training loaders draw indices through the configured MixingStrategy; evaluation loaders (training=False) iterate sequentially for a stable, full-coverage pass.

Parameters:

datasets – Source datasets, or None to build them from datasets.
schemas – Per-dataset schemas (required with datasets).
training – Whether this is the training split.
epoch – Initial epoch for the sampler and per-sample augmentation.

Returns:

A configured torch.utils.data.DataLoader.