blinklinmult.train.cli

Entry point for blink training runs.

Example

Train the joint model over every corpus:

make train-joint

Or directly, for control over the individual configs:

uv run python -m blinklinmult.train.cli \
    --data config/data/all.yaml \
    --model config/model/blinklinmult.yaml \
    --train config/train/joint.yaml

Override any config value without editing a file. The mixing temperature is a dataloader knob, so this needs no rebuild:

make train-joint ARGS="--set data.strategy_kwargs.temperature=3.0"

Attributes

logger

Module-level logger.

MONITOR

mean F1 across supervised targets.

DEFAULT_DATASET_CONFIG_DIR

Directory of per-corpus declarations, resolved by name from the data config.

DEFAULT_REPORT_FPS

Frame rate assumed when no corpus in the run declares one.

MAX_EPOCHS_FILE

Name of the file recording a run's planned max_epochs.

Functions

run_output_dir(→ pathlib.Path)

Directory this run writes its outputs into.

report_fps(→ float)

Frame rate for turning frame counts into minutes.

build_callbacks(→ list[lightning.Callback])

Assemble a run's callbacks.

check_max_epochs(→ None)

Reject a resume that changes the length of the LR schedule.

resume_from(→ str | None)

The checkpoint a run should continue from.

unfreeze_module(...)

Load a frozen-stage run and continue it with the encoder trainable.

embedding_key_for(→ blinklinmult.data.embeddings.CacheKey)

Identity of the embedding cache this run may read.

run(→ dict[str, Any])

Train, then test, one configured experiment.

main(→ None)

CLI entry point.

Module Contents

blinklinmult.train.cli.logger

Module-level logger.

blinklinmult.train.cli.MONITOR = 'valid/mean_f1'

mean F1 across supervised targets.

Type:

Metric that selects the best checkpoint

blinklinmult.train.cli.DEFAULT_DATASET_CONFIG_DIR

Directory of per-corpus declarations, resolved by name from the data config.

blinklinmult.train.cli.run_output_dir(config: blinklinmult.train.config.ExperimentConfig) → pathlib.Path

Directory this run writes its outputs into.

<output_dir>/<experiment>/<run_name>, falling back to the experiment directory when the run is unnamed.

Grouped by experiment, not by task, so that a sweep’s runs sit together on disk exactly as they do in MLflow – one directory per mlflow.experiment_name, matching the experiment the same runs are logged under. Grouping by task instead put every eye-state run ever trained into one directory, which is how a backbone sweep’s checkpoints ended up as best-v25.ckpt .. best-v28.ckpt with nothing to say which backbone produced which.

Named runs get their own subdirectory because every run writes the same filenames – metrics_test.json, test_events.json, test_predictions_<target>.csv, checkpoints/best.ckpt – so a sweep sharing one directory leaves only the last run’s results, which reads as a finished sweep rather than as an error.

Parameters:

config (ExperimentConfig) – The run’s configuration.

Returns:

Directory for this run’s outputs.

Return type:

Path

blinklinmult.train.cli.DEFAULT_REPORT_FPS = 30.0

Frame rate assumed when no corpus in the run declares one.

blinklinmult.train.cli.report_fps(specs: list[blinklinmult.data.schema.DatasetSpec] | None) → float

Frame rate for turning frame counts into minutes.

Parameters:

specs (list[DatasetSpec] | None) – The run’s corpora.

Returns:

The highest rate any of them declares, or DEFAULT_REPORT_FPS when none does.

Return type:

float

blinklinmult.train.cli.build_callbacks(config: blinklinmult.train.config.ExperimentConfig, output_dir: pathlib.Path, specs: list[blinklinmult.data.schema.DatasetSpec] | None = None, eval_only: bool = False, resumable: bool = False, shared_threshold_dir: pathlib.Path | None = None) → list[lightning.Callback]

Assemble a run’s callbacks.

Parameters:
  • config (ExperimentConfig) – The run’s configuration.

  • output_dir (Path) – Directory for run outputs.

  • specs (list[DatasetSpec] | None) – The run’s resolved corpora, for the frame rate the event report needs.

  • shared_threshold_dir (Path | None) – Read and write the event threshold cache here rather than in the run’s own directory, so a per-corpus evaluation shares one universal operating point.

  • resumable (bool) – Periodically checkpoint the test pass so a killed run can resume it. Off by default: the state file is as large as the accumulated predictions, and only a very long pass justifies it.

  • eval_only (bool) – Keep only the callbacks that score and dump the test split. Training-only callbacks – checkpoints, early stopping, the LR monitor, epoch propagation, and timing – are omitted so that an --eval-only run re-fits the event threshold and writes the report artifacts without writing a checkpoint or recording timings.

Returns:

The callbacks.

Return type:

list[L.Callback]

blinklinmult.train.cli.MAX_EPOCHS_FILE = 'max_epochs.txt'

Name of the file recording a run’s planned max_epochs.

Cosine annealing is a fixed-length curve built at first launch for that total. Resuming with a different max_epochs would silently change the LR plan mid-run – the schedule would anneal to the wrong floor, and the only trace would be a loss curve that looks slightly off. Recorded here so the mismatch is rejected instead.

blinklinmult.train.cli.check_max_epochs(output_dir: pathlib.Path, max_epochs: int, resume: bool) → None

Reject a resume that changes the length of the LR schedule.

Parameters:
  • output_dir (Path) – The run’s output directory.

  • max_epochs (int) – What the config asks for now.

  • resume (bool) – Whether this launch is a resume.

Raises:

ValueError – If a resumed run’s max_epochs differs from the original.

blinklinmult.train.cli.resume_from(resume: pathlib.Path | bool | None, output_dir: pathlib.Path) → str | None

The checkpoint a run should continue from.

Parameters:
  • resume (Path | bool | None) – True to continue this run’s own last.ckpt, a path to continue a specific checkpoint, or None/False to start fresh.

  • output_dir (Path) – The run’s output directory.

Returns:

The checkpoint path, or None to start at epoch 0.

Return type:

str | None

Raises:

FileNotFoundError – If a resume was asked for and no checkpoint exists. Silently starting from scratch would waste the hours the run was meant to preserve, and the loss curve is the only place it would show.

blinklinmult.train.cli.unfreeze_module(checkpoint: pathlib.Path, config: blinklinmult.train.config.ExperimentConfig, image_size: int, eye_feature_dim: int | None, steps_per_epoch: int) → blinklinmult.train.module.BlinkLightningModule

Load a frozen-stage run and continue it with the encoder trainable.

The two-stage recipe: train the transformer against a fixed encoder, which is 7.7x cheaper and cannot damage the pretrained features, then refine everything jointly at a much lower rate.

Neither existing path does this. model.encoder_weights loads only the encoder, discarding the transformer that stage 1 spent its whole run training. --resume keeps the whole model but also restores the stored encoder_freeze, so it continues stage 1 rather than starting stage 2.

A fresh optimiser is deliberate. Stage 1’s Adam moments describe a parameter set that excluded the encoder, and its schedule is mid-anneal for a different problem; carrying either forward would apply stage 1’s momentum to weights it never saw.

Parameters:
  • checkpoint (Path) – Stage 1’s checkpoint, usually its best.ckpt.

  • config (ExperimentConfig) – The stage 2 configuration, supplying the low learning rates.

  • image_size (int) – Crop size, from the datasets.

  • eye_feature_dim (int | None) – Feature width, from the datasets.

  • steps_per_epoch (int) – Optimizer steps per epoch.

Returns:

The stage 1 model, fully trainable.

Return type:

BlinkLightningModule

Raises:

FileNotFoundError – If the checkpoint does not exist. Silently starting from scratch would discard the stage it was meant to continue, and only the loss curve would show it.

blinklinmult.train.cli.embedding_key_for(config: blinklinmult.train.config.ExperimentConfig) → blinklinmult.data.embeddings.CacheKey

Identity of the embedding cache this run may read.

Built from a throwaway encoder loaded from the run’s own model.encoder_weights, because the cache is keyed by the weights that produced it and the datamodule is constructed before the model. Loading the encoder twice costs a few seconds once per run and removes any chance of the key describing different weights than the run actually uses.

Parameters:

config (ExperimentConfig) – The resolved run configuration.

Returns:

What a cache shard must match to be read.

Return type:

CacheKey

blinklinmult.train.cli.run(config: blinklinmult.train.config.ExperimentConfig, root: pathlib.Path = PROJECT_ROOT, config_dir: pathlib.Path | None = None, fast_dev_run: bool = False, resume: pathlib.Path | bool | None = None, eval_only: pathlib.Path | None = None, unfreeze_from: pathlib.Path | None = None, reuse_threshold: bool = False, resumable: bool = False, shared_threshold_dir: pathlib.Path | None = None) → dict[str, Any]

Train, then test, one configured experiment.

Parameters:
  • config (ExperimentConfig) – The run’s configuration.

  • root (Path) – Repository root that data paths resolve against.

  • config_dir (Path | None) – Directory of per-corpus declarations. Defaults to config/data.

  • fast_dev_run (bool) – Run a single batch through every stage, for smoke-testing.

  • resume (Path | bool | None) – True continues this run from its own last.ckpt and appends to its MLflow run; a path continues a specific checkpoint; None starts at epoch 0.

  • unfreeze_from (Path | None) – Stage-1 checkpoint to continue with the encoder trainable, or None for a single-stage run. See unfreeze_module() for why neither model.encoder_weights nor --resume does this.

  • eval_only (Path | None) – Checkpoint to score on the test split without training. The module is loaded from this file, the event operating point is fitted on the validation split, and the test split is then scored – the same passes a finished run performs, so its artifacts are directly comparable to one. Use it to recover the benchmark of a run that was interrupted before its test pass.

  • reuse_threshold (bool) – Skip the validation pass when a cached event threshold from a previous run is available. Fitting the operating point is the only thing validation contributes to an eval-only run, and it costs a full pass over the split, so a retry after an interrupted test pass can reuse the cached value instead. Falls back to fitting when no usable cache exists.

  • resumable (bool) – Checkpoint the test pass periodically so a killed run can pick it up. The partial state is deleted once the pass completes, so it costs nothing on disk afterwards.

  • shared_threshold_dir (Path | None) – Share the fitted event threshold through this directory. A per-corpus evaluation points every run at one location, so the corpora that can fit a threshold write it and the rest – including those with no validation split – read the same value instead of falling back to a fixed default.

Returns:

The test metrics.

Return type:

dict[str, float]

blinklinmult.train.cli.main() → None

CLI entry point.