nvalchemi.training.OptimizerConfig

pydantic model nvalchemi.training.OptimizerConfig

Declarative optimizer and optional LR-scheduler bundle.

OptimizerConfig captures how to build a torch optimizer (and an optional learning-rate scheduler) without instantiating either until model parameters are available. You supply the optimizer class plus its keyword arguments as data, so the whole training recipe stays serializable and can round-trip through to_spec() / from_spec(). This is the unit that TrainingStrategy and FineTuningStrategy accept via their optimizer_configs argument, either as a single config for a single-model strategy or as a {model_name: [OptimizerConfig, ...]} mapping for named/multi-model strategies.

Both optimizer_kwargs and scheduler_kwargs are validated against the corresponding class __init__ signature at construction time (see _validate_kwargs), so a misspelled or unsupported argument raises a pydantic.ValidationError immediately rather than deep inside the training loop. At runtime, build() binds the config to a concrete parameter iterable and returns the (optimizer, scheduler) pair; the strategy calls this for you once trainable parameters are known.

Schedulers fall into two families that the strategy steps at different times. Time-based schedulers (StepLR, CosineAnnealingLR, etc.) advance on every optimizer step via step_lr_schedulers(). Metric-driven schedulers (ReduceLROnPlateau and subclasses) step only at validation checkpoints via step_metric_schedulers(), using scheduler_metric_adapter to pull a scalar out of the validation summary dict.

Examples

A plain Adam optimizer with no scheduler:

>>> import torch
>>> cfg = OptimizerConfig(
...     optimizer_cls=torch.optim.Adam,
...     optimizer_kwargs={"lr": 1e-3},
... )

A time-based StepLR schedule that decays every 10 optimizer steps:

>>> cfg = OptimizerConfig(
...     optimizer_cls=torch.optim.Adam,
...     optimizer_kwargs={"lr": 1e-3},
...     scheduler_cls=torch.optim.lr_scheduler.StepLR,
...     scheduler_kwargs={"step_size": 10, "gamma": 0.1},
... )

A metric-driven ReduceLROnPlateau schedule keyed off a validation summary entry (stepped only at validation checkpoints):

>>> cfg = OptimizerConfig(
...     optimizer_cls=torch.optim.AdamW,
...     optimizer_kwargs={"lr": 1e-4},
...     scheduler_cls=torch.optim.lr_scheduler.ReduceLROnPlateau,
...     scheduler_kwargs={"mode": "min", "patience": 5},
...     scheduler_metric_adapter="total_loss",
... )

Attach configs to a strategy — a bare config for a single model, or a per-model mapping for named models:

>>> from nvalchemi.training import TrainingStrategy
>>> strategy = TrainingStrategy(
...     models=model,
...     optimizer_configs=cfg,
...     num_epochs=10,
...     training_fn=default_training_fn,
...     loss_fn=EnergyMSELoss(),
... )

Notes

• scheduler_kwargs must be empty unless scheduler_cls is set, and scheduler_metric_adapter may only be supplied alongside a scheduler_cls; violating either raises at construction time.

• scheduler_metric_adapter is only consulted for metric-driven schedulers. A str is a key lookup into the validation summary, a callable receives the whole summary and returns a float, and None falls back to the default "total_loss" key.

• Direct torch.nn.Module/class references are accepted at runtime, but only importable classes serialize cleanly through to_spec().

field optimizer_cls: type [Required]

Optimizer class; optimizer_kwargs must match its signature.

Constraints:

• func = <function _serialize_type at 0xeb0b2465fb00>

• json_schema_input_type = PydanticUndefined

• return_type = PydanticUndefined

• when_used = always

field optimizer_kwargs: dict[str, Any] [Optional]

Keyword arguments forwarded to the optimizer constructor; validated against its __init__ signature at construction time.

field scheduler_cls: type | None = None

Optional LR scheduler. Time-based schedulers (StepLR, CosineAnnealingLR, etc.) step every optimizer step. Metric-driven schedulers (ReduceLROnPlateau and subclasses) step only at validation checkpoints via step_metric_schedulers().

Constraints:

• func = <function _serialize_type at 0xeb0b2465fb00>

• json_schema_input_type = PydanticUndefined

• return_type = PydanticUndefined

• when_used = always

field scheduler_kwargs: dict[str, Any] [Optional]

Must be empty unless scheduler_cls is set.

field scheduler_metric_adapter: Callable[[dict[str, Any]], float] | str | None = None

How a metric-driven scheduler (ReduceLROnPlateau) extracts its scalar metric from the validation summary dict. A str is treated as a key lookup into the summary; a callable receives the whole summary dict and returns a float; None uses the default extractor (see _extract_scheduler_metric()).