"""
colibri.optax_optimizer.py
Module contains functions for optax gradient descent optimisation.
"""
import optax
from flax.training.early_stopping import EarlyStopping
import logging
log = logging.getLogger(__name__)
[docs]
def optimizer_provider(
optimizer_settings,
) -> optax._src.base.GradientTransformationExtraArgs:
"""
Define the optimizer.
Parameters
----------
optimizer_settings : dict, default = {}
Dictionary containing the optimizer settings.
Returns
-------
optax._src.base.GradientTransformationExtraArgs
Optax optimizer.
"""
optimizer = optimizer_settings["optimizer"]
log.info(f"Using {optimizer} optimizer.")
opt = getattr(optax, optimizer)
optimizer_hyperparams = optimizer_settings.get("optimizer_hyperparams", {})
log.info(f"Optimizer hyperparameters: {optimizer_hyperparams}.")
scheduler_cfg = optimizer_settings.get("scheduler")
if scheduler_cfg is not None:
scheduler = scheduler_cfg["name"]
sched_params = scheduler_cfg["params"]
log.info(f"Using {scheduler} scheduler with params {sched_params}.")
scheduler_fn = getattr(optax, scheduler)(**sched_params)
optimizer_hyperparams["learning_rate"] = scheduler_fn
clipnorm_cfg = optimizer_settings.get("clipnorm")
if clipnorm_cfg is not None:
clipnorm = optimizer_settings["clipnorm"]
log.info(f"Using gradient clipping with norm {clipnorm}.")
return optax.chain(
optax.clip_by_global_norm(clipnorm),
opt(**optimizer_hyperparams),
)
else:
return opt(**optimizer_hyperparams)
[docs]
def early_stopper(
min_delta=1e-5, patience=20, max_epochs=1000, mc_validation_fraction=0.2
):
"""
Define the early stopping criteria.
When ``mc_validation_fraction`` is zero, the full training loss is used as
the monitoring loss, as in n3fit. The configured patience therefore still
applies. ``max_epochs`` and ``mc_validation_fraction`` are retained in the
signature for runcard compatibility.
"""
return EarlyStopping(min_delta=min_delta, patience=patience)