nvalchemi.training.setup_optimizers

nvalchemi.training.setup_optimizers(models, optimizer_configs, *, allowed_parameter_names=None)

Build optimizers and schedulers for configured model names.

Parameters:

• models (torch.nn.Module | dict[str, torch.nn.Module] | torch.nn.ModuleDict)

• optimizer_configs (OptimizerConfig | list[OptimizerConfig] | dict[str, list[OptimizerConfig]])

• allowed_parameter_names (set[str] | None, optional) – Fully qualified parameter names allowed in optimizers. Names use the format "<model_key>.<parameter_name>". When None all requires_grad=True parameters are eligible.

Return type:

dict[str, list[tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler | None]]]

Raises:

ValueError – If a config key is not present in models or a configured model has no trainable parameters.