blinklinmult.train.module

The Lightning module wrapping BlinkModel.

One training step serves all three model families and both tasks. What makes that possible is that every target is handled identically — logits, target, and mask — so a joint run and a single-task run differ only in how many entries the target loop has, and a corpus that does not annotate a target simply contributes an all-False mask that the loss and the metrics both skip.

Attributes

logger

Module-level logger.

Classes

BlinkLightningModule

Trains the blink model on one or both tasks.

Module Contents

blinklinmult.train.module.logger

Module-level logger.

class blinklinmult.train.module.BlinkLightningModule(model_config: blinklinmult.train.config.ModelConfig | collections.abc.Mapping[str, Any], train_config: blinklinmult.train.config.TrainConfig | collections.abc.Mapping[str, Any], image_size: int, eye_feature_dim: int | None = None, steps_per_epoch: int = 1)

Bases: lightning.LightningModule

Trains the blink model on one or both tasks.

The model is built here from its config rather than injected, so load_from_checkpoint(path) reconstructs a run without the caller having to rebuild and pass the architecture.

Parameters:

  • model_config (ModelConfig) – Architecture settings.

  • train_config (TrainConfig) – Task, loss, and optimizer settings.

  • image_size (int) – Side length of a square eye crop, from the datasets.

  • eye_feature_dim (int | None) – Handcrafted feature width, from the datasets, or None when the run has none.

  • steps_per_epoch (int) – Optimizer steps per epoch, for step schedules.

property feature_names: list[str]

Feature keys this model reads from a batch, in input order.

Returns:

One or two keys, depending on the model family.

Return type:

list[str]

forward(batch: dict[str, Any]) → dict[str, torch.Tensor]

Run the model over one batch.

Parameters:

batch (dict) – A collated OmniLoader batch.

Returns:

Raw logits per target.

Return type:

dict[str, torch.Tensor]

training_step(batch: dict[str, Any], batch_idx: int) → torch.Tensor

Run one training step.

Parameters:

  • batch (dict) – A collated batch.

  • batch_idx (int) – Batch index within the epoch.

Returns:

The loss to backpropagate.

Return type:

torch.Tensor

validation_step(batch: dict[str, Any], batch_idx: int) → torch.Tensor

Run one validation step.

Parameters:

  • batch (dict) – A collated batch.

  • batch_idx (int) – Batch index within the epoch.

Returns:

The loss.

Return type:

torch.Tensor

test_step(batch: dict[str, Any], batch_idx: int) → torch.Tensor

Run one test step, keeping sample provenance for the artifact dump.

Parameters:

  • batch (dict) – A collated batch.

  • batch_idx (int) – Batch index within the epoch.

Returns:

The loss.

Return type:

torch.Tensor

predict_step(batch: dict[str, Any], batch_idx: int) → dict[str, Any]

Predict one batch, keeping sample ids alongside the probabilities.

Parameters:

  • batch (dict) – A collated batch.

  • batch_idx (int) – Batch index.

Returns:

key, dataset, and a probability per target.

Return type:

dict

on_train_epoch_end() → None

Log the training epoch’s metrics.

on_validation_epoch_start() → None

Clear the per-sample ids before a validation epoch.

Cleared at the start rather than the end so a callback reading them after the epoch still sees the epoch that just ran.

on_validation_epoch_end() → None

Log the validation epoch’s metrics.

release_validation_state() → None

Free everything the validation epoch accumulated.

Validation state is deliberately kept alive past its own epoch so the callbacks that fit the event operating point can read it (see on_validation_epoch_start()). In a fit run the next epoch clears it; in an eval-only run nothing does, so the whole validation split stays resident through the test pass – measured at 4.75 GB of accelerator memory already held at test batch 0.

Called between validate() and test() once the threshold has been fitted, so it frees the state without changing any reported number.

on_test_epoch_start() → None

Clear per-sample and metric state before a test epoch.

Resetting here rather than at epoch end keeps the accumulated predictions readable by the artifact-writing callbacks, whatever order they run in, while still keeping repeated trainer.test() calls from accumulating on top of each other.

on_test_epoch_end() → None

Log the test epoch’s metrics.

The accumulator is deliberately left populated: the prediction-writing and plotting callbacks read it after this hook.

property lr: float

The optimizer’s peak learning rate.

A thin proxy onto train_config.optimizer.lr under the flat name Lightning’s LR finder requires: it looks for a top-level lr or learning_rate attribute to read, and writes the swept value back through it. See lr_find.

Nothing in the normal training path reads this – configure_optimizers() reads train_config.optimizer.lr directly, so the proxy cannot drift from what actually trains.

Returns:

The peak learning rate for the transformer and head group. backbone_lr is a separate, deliberately lower rate and is not touched by the finder.

Return type:

float

configure_optimizers() → dict[str, Any]

Build the optimizer and LR schedule.

Returns:

Lightning’s optimizer configuration.

Return type:

dict