1 lightning.one_step.datamodule

lightning.one_step.datamodule

Provide stage-aware VIN datasets and candidate-table collation to Lightning.

The training data-flow mirrors aria_nbv/scripts/train_vin.py:

EFM snippet → candidate generation → depth rendering → backprojection → oracle RRI → VIN (CORAL).

This module owns dataset construction, stage reuse, candidate shuffling, and VinOracleBatch.collate selection. It keeps the expensive oracle labeler in the data pipeline by default, but can switch to immutable VIN offline stores for fast parallel reading. It does not own scorer inputs or losses; those boundaries belong to :mod:aria_nbv.lightning.one_step._batch and aria_nbv.lightning.one_step.module.

1.1 Attributes

Name	Description
VinDatasetSourceConfig	Split-aware VIN dataset-source union composed by Lightning.

1.2 Classes

Name	Description
VinDataModule	Own stage datasets and loaders for oracle-labelled VIN batches.
VinDataModuleConfig	Configure stage datasets and candidate-table collation.

1.2.1 VinDataModule

lightning.one_step.datamodule.VinDataModule(config)

Own stage datasets and loaders for oracle-labelled VIN batches.

Lightning may call :meth:setup more than once. Dataset instances are therefore initialized lazily and retained by stage. Scorer-device transfer remains the module’s responsibility; this class yields CPU/source-device VinOracleBatch objects whose candidate width is compact for a single sample or right-padded after offline collation.

1.2.1.1 Attributes

Name	Type	Description
config		Immutable construction policy for sources, batching, and worker lifecycle.

1.2.1.2 Methods

Name	Description
setup	Initialize and retain the datasets required by one Lightning stage.
train_dataloader	Build the training loader with optional row and candidate shuffling.
val_dataloader	Build the deterministic validation loader for the retained source.
test_dataloader	Reuse the validation-loader contract for Lightning test evaluation.
iter_oracle_batches	Return a direct iterator over one stage’s oracle-labelled samples.

1.2.1.2.1 setup

lightning.one_step.datamodule.VinDataModule.setup(stage=None)

Initialize and retain the datasets required by one Lightning stage.

1.2.1.2.1.1 Parameters

Name	Type	Description	Default
stage	Stage | str | None	train, val, test, or None. None prepares both training and validation ownership paths.	None

1.2.1.2.1.2 Notes

Repeated calls are idempotent for already-created sources. Validation and test share _val_source; when use_train_as_val is enabled they deliberately share _train_source as well.

1.2.1.2.2 train_dataloader

lightning.one_step.datamodule.VinDataModule.train_dataloader()

Build the training loader with optional row and candidate shuffling.

1.2.1.2.2.1 Returns

Name	Type	Description
	DataLoader	DataLoader yielding VinOracleBatch. Map-style batching
	DataLoader	produces poses PoseTW["B N_q 12"], RRI labels
	DataLoader	Tensor["B N_q", float32], and a hard prefix mask derived from
	DataLoader	candidate_count; iterable/online sources remain unbatched.

1.2.1.2.3 val_dataloader

lightning.one_step.datamodule.VinDataModule.val_dataloader()

Build the deterministic validation loader for the retained source.

1.2.1.2.3.1 Returns

Name	Type	Description
	DataLoader	DataLoader yielding unshuffled VinOracleBatch objects with
	DataLoader	the same compact/right-padded shape contract as the training loader.

1.2.1.2.4 test_dataloader

lightning.one_step.datamodule.VinDataModule.test_dataloader()

Reuse the validation-loader contract for Lightning test evaluation.

1.2.1.2.5 iter_oracle_batches

lightning.one_step.datamodule.VinDataModule.iter_oracle_batches(stage)

Return a direct iterator over one stage’s oracle-labelled samples.

This bypasses worker processes and VinOracleBatch.collate, so yielded candidate payloads retain their per-sample N_q width.

1.2.2 VinDataModuleConfig

lightning.one_step.datamodule.VinDataModuleConfig()

Configure stage datasets and candidate-table collation.

Online sources yield one oracle-labelled snippet at a time and therefore require num_workers=0. Map-style offline sources may collate a batch of size B; candidate rows are right-padded to N_q and accompanied by per-sample candidate_count values used to construct the hard prefix mask.

1.2.2.1 Attributes

Name	Description
target_type	Return the VinDataModule factory target.
paths	Project path resolver used by nested dataset/store configurations.
source	Config-as-factory source for online labels or an immutable offline store.
shuffle	Shuffle map-style training samples each epoch; ignored for iterable sources.
shuffle_candidates	Permute compact valid candidate rows and aligned labels per training sample.
num_workers	DataLoader worker count; must be zero for the non-multiprocess-safe online labeler.
batch_size	Optional offline batch size B using VinOracleBatch.collate.
persistent_workers	Keep worker processes between epochs; forced off when num_workers is zero.
verbosity	Dataset and oracle-labeler diagnostic verbosity.
is_debug	Serialized compatibility flag; worker and trainer debug policy is configured elsewhere.
use_train_as_val	Reuse the exact train dataset instance for validation and testing.