omniloader.integrations.lightning

Optional PyTorch Lightning integration.

This module wires OmniLoaders and mixing strategies into a LightningDataModule. Lightning is an optional dependency: importing this module without it installed raises a clear error, while the rest of OmniLoader keeps working with plain torch.utils.data loaders. Install via the omniloader[lightning] extra.

Classes

OmniDataModule

A LightningDataModule that mixes datasets per split via OmniLoader.

Module Contents

class omniloader.integrations.lightning.OmniDataModule(config: omniloader.config.OmniConfig, train: tuple[collections.abc.Sequence[omniloader.loader.SizedDataset], collections.abc.Sequence[omniloader.schema.spec.DatasetSchema]] | None = None, valid: tuple[collections.abc.Sequence[omniloader.loader.SizedDataset], collections.abc.Sequence[omniloader.schema.spec.DatasetSchema]] | None = None, test: tuple[collections.abc.Sequence[omniloader.loader.SizedDataset], collections.abc.Sequence[omniloader.schema.spec.DatasetSchema]] | None = None, schema: omniloader.schema.spec.UnifiedSchema | None = None)

Bases: _LightningBase

A LightningDataModule that mixes datasets per split via OmniLoader.

The train split is mixed with the configured strategy through an OmniSampler; validation and test splits are iterated in order (proportional, no shuffling) for stable evaluation.

Parameters:

  • config – The run configuration (seed, batch size, strategy, …).

  • train – (datasets, schemas) for the training split, or None.

  • valid – (datasets, schemas) for the validation split, or None.

  • test – (datasets, schemas) for the test split, or None.

  • schema – Optional shared unified schema applied to every split. When None, each split unions its own schemas.

setup(stage: str | None = None) → None

Instantiate the loaders and the training sampler.

Parameters:

stage – Lightning stage hint (unused; all splits are built lazily).

set_epoch(epoch: int) → None

Propagate the epoch to every loader and the training sampler.

Ensures both the mixing draw and per-sample augmentation advance each epoch while staying reproducible.

Parameters:

epoch – The upcoming epoch number.

train_dataloader() → torch.utils.data.DataLoader | None

Return the training DataLoader (or None if no train split).

val_dataloader() → torch.utils.data.DataLoader | None

Return the validation DataLoader (or None).

test_dataloader() → torch.utils.data.DataLoader | None

Return the test DataLoader (or None).