blinklinmult.train.optim¶

Optimizer and learning-rate schedule construction.

The one thing here that is specific to this project is the discriminative learning rate: an ImageNet-pretrained eye backbone and a randomly-initialised transformer do not want the same step size. Fine-tuning the backbone at the transformer’s rate destroys the pretrained features in the first few hundred steps, which shows up as a run that trains stably and scores worse than the frozen-backbone baseline. build_optimizer() therefore splits parameters into two groups whenever optimizer.backbone_lr is set.

Attributes¶

`logger`	Module-level logger.
`BACKBONE_GROUP`	Name of the parameter group holding the pretrained image backbone.
`HEAD_GROUP`	Name of the parameter group holding everything else.
`MIN_ONECYCLE_STEPS`	Fewest optimizer steps a one-cycle schedule can be given.

Exceptions¶

OptimError

Raised when an optimizer or schedule cannot be built.

Functions¶

`parameter_groups`(→ list[dict[str, Any]])	Split a model's parameters into backbone and head groups.
`build_optimizer`(→ torch.optim.Optimizer)	Construct the configured optimizer over the model's parameter groups.
`build_scheduler`(→ dict[str, Any] \| None)	Construct the configured LR schedule, in Lightning's config format.

Module Contents¶

blinklinmult.train.optim.logger¶: Module-level logger.

blinklinmult.train.optim.BACKBONE_GROUP = 'backbone'¶: Name of the parameter group holding the pretrained image backbone.

blinklinmult.train.optim.HEAD_GROUP = 'head'¶: Name of the parameter group holding everything else.

blinklinmult.train.optim.MIN_ONECYCLE_STEPS = 3¶

Fewest optimizer steps a one-cycle schedule can be given.

Below this there is no room for a warmup step, a peak, and an annealing step, and OneCycleLR divides by a zero-width phase.

exception blinklinmult.train.optim.OptimError¶

Bases: ValueError

Raised when an optimizer or schedule cannot be built.

Initialize self. See help(type(self)) for accurate signature.

blinklinmult.train.optim.parameter_groups(model: torch.nn.Module, config: blinklinmult.train.config.OptimConfig) → list[dict[str, Any]]¶

Split a model’s parameters into backbone and head groups.

Parameters:

model (nn.Module) – The model.
config (OptimConfig) – Optimizer settings.

Returns:

Optimizer parameter groups. A single group when backbone_lr is unset, or when the backbone is frozen and therefore contributes no trainable parameter.

Return type:

list[dict]

blinklinmult.train.optim.build_optimizer(config: blinklinmult.train.config.OptimConfig, model: torch.nn.Module) → torch.optim.Optimizer¶

Construct the configured optimizer over the model’s parameter groups.

Parameters:

config (OptimConfig) – Optimizer settings.
model (nn.Module) – The model whose parameters are optimized.

Returns:

The optimizer.

Return type:

Optimizer

Raises:

OptimError – If the optimizer name is unknown, or the model has no trainable parameter at all.

blinklinmult.train.optim.build_scheduler(config: blinklinmult.train.config.OptimConfig, optimizer: torch.optim.Optimizer, steps_per_epoch: int, max_epochs: int) → dict[str, Any] | None¶

Construct the configured LR schedule, in Lightning’s config format.

Parameters:

config (OptimConfig) – Optimizer settings.
optimizer (Optimizer) – The optimizer to schedule.
steps_per_epoch (int) – Optimizer steps in one epoch.
max_epochs (int) – Total epochs.

Returns:

Lightning lr_scheduler_config, or None when no schedule is configured.

Return type:

dict | None

Raises:

OptimError – If the scheduler name is unknown.