"One loader to load them all, one schema to find them, one batch to bring them all, and in the mask bind them."

Problem definition

Multi-task learning wants one model to learn many related tasks at once. In practice the supervision for those tasks is scattered across separate datasets that annotate completely different things: one corpus has a single label per sample (say, sentiment), another has a per-step label sequence (valence/arousal), a third has categorical class ids — and each one only covers its own slice of the features. You cannot just concatenate them and call DataLoader; the samples do not even share a shape.

OmniLoader builds the union of every feature and target across your datasets and yields each sample in that one shared format. It copies the keys a dataset provides, pads or crops sequences to the lengths you declare, and fills in the keys a dataset is missing with placeholder tensors plus an all-False mask — so your model always knows what was real and what was absent. It is modality-agnostic: video, audio, text, or anything else, as long as it is vectors and sequences.

Tasks it solves

What's in the box

Category	Components	Description
Schema & unification	OmniLoader, SampleUnifier, TensorSpec, DatasetSchema, UnifiedSchema, unified_collate	Declare each dataset's tensors and build the union schema every sample is mapped into, filling missing keys with placeholders and masks.
Dataset loading	HDF5Dataset, NpyFolderDataset, DictTensorDataset, split_indices	Read your features from disk or memory in the usual layouts, and carve reproducible train/val/test splits.
Mixing strategies	ProportionalStrategy, TemperatureStrategy, AnnealedTemperatureStrategy, FixedWeightStrategy, RoundRobinStrategy	Decide how often each dataset is sampled so large corpora do not drown out small ones during joint training.
Class balance	class_weights_for_sampler, class_weights_for_loss, class_histogram	Inspect label distributions and derive weights to counter class imbalance on either the sampler or the loss side.
Normalization	Normalize, MinMax, Robust, Instance, PerDatasetNormalize, compute_stats, compute_dataset_stats	Scale features with global or per-dataset statistics to remove annotation-source leakage across corpora.
Augmentation	GaussianNoise, FeatureDropout, SpanMasking, FeatureMasking, TimeWarp, RandomCrop, CenterCrop, MixupCollator, Compose	Perturb features and sequences to regularize training and build tolerance to missing or corrupted inputs.
Batching	DynamicCollator, LengthBucketBatchSampler	Group variable-length samples and bucket by length to keep batches padding-efficient.
Distributed	OmniSampler, OmniDataModule (Lightning integration)	Scale the unified stream across multiple GPUs and plug straight into a PyTorch Lightning training loop.
Config & CLI	OmniConfig, build_dataloader(), omniloader CLI commands	Specify an entire data pipeline from a config file and build it in code or from the command line for reproducible runs.

Try it

With uv (recommended)

uv add omniloader
uv add "omniloader[lightning]"

Or with pip

pip install omniloader
pip install "omniloader[lightning]"

Requirements: Python 3.12+, PyTorch 2.12.1+, numpy, h5py, pyyaml.