OmniLoader¶

OmniLoader is a PyTorch meta data loader that unifies several datasets with disjoint annotations into a single, masked sample scheme so a model can be trained jointly across all of them.

The problem. Multi-task learning wants one model to learn many related tasks at once, but the supervision for those tasks is usually scattered across separate datasets that annotate different things: one corpus has a single label per sample (e.g. sentiment), another a per-step label sequence (e.g. valence/arousal), a third categorical class ids — and each covers only its own subset of the features. No single dataset carries every task’s labels, so a shared-backbone model cannot simply be pointed at one file. The samples first have to be brought to one common, batchable scheme where every task’s slot is always present and the loss can tell a real label apart from a missing one.

What OmniLoader does. It builds the union of every feature and target across your datasets. Every sample it yields carries the full union schema: keys a dataset provides are copied (and, for sequences, padded/cropped to a declared length), while keys it lacks are filled with a placeholder tensor plus an all-False <name>_mask. So a single batch can freely mix samples from different datasets, each task’s head sees a consistently-shaped tensor, and your loss and model always know which values are real. It is a map-style torch.utils.data.Dataset you wrap in an ordinary DataLoader — it unifies and masks, it does not reinvent batching.

OmniLoader is modality-, dataset- and model-agnostic: it knows nothing about video, audio, text, which specific corpora you loaded, or any model. It reasons only about vectors (shape () or (F,)) and sequences (shape (T,) or (T, F)); a value is a sequence exactly when its spec sets time_dim. A value can also be structured — e.g. an image (C, H, W) — by setting shape instead of feature_dim, yielding (*shape,) or, for a sequence, (T, *shape).

Beyond this core unification and its utilities (schema declaration, dataset adapters, splits, introspection), OmniLoader also ships the surrounding machinery a joint training run needs: mixing strategies to balance datasets of very different sizes, subsampling to control within-dataset draws, class-balance calculations (sampler weights and loss weights), and a mask-aware transform pipeline of normalization and augmentation techniques — all configurable end-to-end from a single file.

Contents

API Reference
- omniloader

Feature overview¶

Category	Feature	What it does
Core	`OmniLoader`	Concatenates disjoint datasets into one masked, unified stream.
	`SampleUnifier`	Maps a raw sample onto the union schema; fills gaps with placeholder + `<name>_mask`.
	`unified_collate()`	Stacks tensors, lists metadata.
Schema	`TensorSpec`	Declare a value (feature or target): `feature_dim`/`shape`, `time_dim`, dtype, placeholder.
	`DatasetSchema` / `UnifiedSchema`	Per-dataset specs; merged + validated union.
	vector / sequence	Structural, modality-agnostic (`()`, `(F,)`, `(T,)`, `(T, F)`).
Datasets	`HDF5Dataset`	Per-sample HDF5 groups; worker-safe; `cache_size` / `preload`.
	`NpyFolderDataset`	Memory-mapped `root/<sample>/<key>.npy` layout.
	`DictTensorDataset`	In-memory random tensors (tests/experiments).
	`split_indices()`	Reproducible, optionally stratified train/val/test splits.
Mixing	`ProportionalStrategy`	Sample by true dataset sizes.
	`TemperatureStrategy`	`size ** (1 / T)` re-weighting (T=2 → sqrt).
	`AnnealedTemperatureStrategy`	Temperature annealed per epoch.
	`FixedWeightStrategy`	Explicit or uniform per-dataset weights.
	`RoundRobinStrategy`	Equal, interleaved draws.
Subsampling	`SubsampleConfig` / `IndexPool`	Replacement, FRESH/EXHAUST, `effective_size`, per-sample weights.
	`class_weights_for_sampler()`	Per-sample inverse-frequency weights for the sampler.
Class balance	`class_weights_for_loss()` / `class_histogram()`	Per-class loss weights + exact counts for `CrossEntropyLoss(weight=)` (persist via the CLI).
Batching	`DynamicCollator`	Pad sequences to per-batch max, not fixed length (native-length mode only).
	`LengthBucketBatchSampler`	Group similar-length sequences to cut padding (needs `pad_features=False`).
Normalization	`Normalize` / `MinMax` / `Robust` / `Instance`	Standardize features (from stats or per-sample).
	`PerDatasetNormalize`	Standardize each sample by its source dataset’s stats, with an inference fallback for unseen sources.
	`compute_stats()` / `compute_dataset_stats()`	Pooled or per-dataset mean/std/min/max/median/iqr.
	`save_stats()` / `load_stats()`	JSON persistence (flat or per-dataset).
Augmentation	`GaussianNoise`	Feature corruption.
	`FeatureDropout`	Drop whole feature streams (modality dropout).
	`SpanMasking` / `FeatureMasking`	Zero time spans / feature bands.
	`TimeWarp`	Random speed perturbation.
	`RandomCrop` / `CenterCrop`	Windowed sequence crop (features + targets aligned).
	`MixupCollator`	MixUp/CutMix at collate (emits `mixup_lambda` + paired targets).
	`Compose`	Chain transforms (mask-aware, seedable, train/eval-gated).
Distributed	`OmniSampler`	DDP index sharding (`num_replicas` / `rank`).
	`set_epoch` + `seed_worker()`	Reproducible, worker-count-independent augmentation.
Introspection	`describe()`	Coverage matrix, valid fractions, class distributions.
	`validate()`	Dry-run check of data vs declared specs.
Config & CLI	`OmniConfig`	One JSON/YAML file: seed, strategy, subsample, transforms, collate, bucketing, DDP & DataLoader knobs, datasets.
	`build_dataloader()`	Assemble a ready `DataLoader` end-to-end from the config alone.
	`build_datasets()`	Construct datasets from a declarative `datasets` section.
	`config_template_path()`	Locate the shipped, fully annotated `config_template.yaml` (every option + defaults).
	`omniloader` CLI	`describe` / `validate` / `compute-stats` / `class-weights-for-loss`.
Integration	`OmniDataModule`	Optional Lightning module (wires strategy, DDP, seeding).

Quickstart¶

import torch
from torch.utils.data import DataLoader
from omniloader import (
    DatasetSchema, DictTensorDataset, TensorSpec,
    OmniLoader, TemperatureStrategy, unified_collate,
)

# Dataset A: a sequence target (per-step) over a feature sequence.
ds_a = DictTensorDataset({"video": torch.randn(40, 16, 32), "valence": torch.randn(40, 16)})
schema_a = DatasetSchema(
    features=[TensorSpec("video", feature_dim=32, time_dim=16)],  # sequence (T, F)
    targets=[TensorSpec("valence", time_dim=16, placeholder=-5.0)],  # sequence (T,)
)

# Dataset B: a scalar target (one per sample), different feature sequence.
ds_b = DictTensorDataset({"audio": torch.randn(200, 24, 8), "sentiment": torch.randn(200)})
schema_b = DatasetSchema(
    features=[TensorSpec("audio", feature_dim=8, time_dim=24)],  # sequence (T, F)
    targets=[TensorSpec("sentiment", placeholder=-5.0)],  # vector scalar ()
)

omni = OmniLoader([ds_a, ds_b], [schema_a, schema_b])
sampler = omni.make_sampler(TemperatureStrategy(omni.dataset_sizes, temperature=2.0))
loader = DataLoader(omni, batch_size=8, sampler=sampler, collate_fn=unified_collate)

for batch in loader:
    # Every batch carries video, audio, valence, sentiment + their masks.
    print(batch["valence"].shape, batch["valence_mask"].shape)
    break

A single dataset works too — pass one-element lists (OmniLoader([ds_a], [schema_a])). The union is then just that dataset’s schema and the multi-dataset machinery (mixing strategies, per-dataset normalization) simply lies dormant.

Config-driven runs¶

A single JSON/YAML file can describe an entire experiment — seed, DataLoader knobs, mixing strategy, per-dataset subsampling, the transform pipeline, the collate function, length bucketing, distributed settings and the datasets themselves. build_dataloader() assembles a ready DataLoader from it:

from omniloader import OmniConfig

config = OmniConfig.from_file("config.yaml")
train_loader = config.build_dataloader(training=True)
val_loader = config.build_dataloader(training=False)  # sequential, augmentations off

A feature that is not a flat vector — e.g. a window of raw image frames — declares shape instead of feature_dim in its spec, either in Python or the same YAML:

TensorSpec(
    name="eye_image",
    time_dim=15,           # T
    shape=(3, 64, 64),     # C, H, W
    dtype=torch.float32,
)
# value_shape == (15, 3, 64, 64); mask_shape == (15,), one flag per frame

OmniLoader¶

Feature overview¶

Quickstart¶

Config-driven runs¶

Indices and tables¶