Added Clipnorm to the MC gradient descent

Author: ecole41

colibri/monte_carlo_fit.py

@@ -16,6 +16,7 @@
 
 from colibri.data_batch import data_batches
 from colibri.mc_utils import write_exportgrid_mc
+from colibri.optax_optimizer import optimizer_provider
 
 log = logging.getLogger(__name__)
 
@@ -52,10 +53,12 @@ def monte_carlo_fit(
     len_trval_data,
     pdf_model,
     mc_initial_parameters,
-    optimizer_provider,
     early_stopper,
     max_epochs,
     FIT_XGRID,
+    optimizer="adam",
+    optimizer_settings={},
+    clipnorm=6.073e-6,
     batch_size=None,
     batch_seed=1,
     alpha=1e-7,
@@ -86,9 +89,6 @@ def monte_carlo_fit(
     mc_initial_parameters: jnp.array
         Initial parameters for the Monte Carlo fit.
 
-    optimizer_provider: optax._src.base.GradientTransformationExtraArgs
-        Optax optimizer.
-
     early_stopper: flax.training.early_stopping.EarlyStopping
         Early stopping criteria.
 
@@ -99,6 +99,15 @@ def monte_carlo_fit(
         xgrid of the theory, computed by a production rule by taking
         the sorted union of the xgrids of the datasets entering the fit.
 
+    optimizer: str, default="adam"
+        The optimizer to use.
+    
+    optimizer_settings: dict, default={}
+        A dictionary with the settings for the optimizer.
+
+    clipnorm: float, optional
+        If not None, gradients will be clipped to this norm.
+
     batch_size: int, default is None which sets it to the full size of data
         Size of batches during training.
 
@@ -121,6 +130,10 @@ def monte_carlo_fit(
 
     pred_and_pdf = pdf_model.pred_and_pdf_func(FIT_XGRID, forward_map=_pred_data)
 
+    optim = optimizer_provider(
+        optimizer=optimizer, optimizer_settings=optimizer_settings, clipnorm=clipnorm
+    )
+
     @jax.jit
     def loss_training(
         parameters,
@@ -173,7 +186,7 @@ def step(
             alpha,
             lambda_positivity,
         )
-        updates, opt_state = optimizer_provider.update(grads, opt_state, params)
+        updates, opt_state = optim.update(grads, opt_state, params)
         params = optax.apply_updates(params, updates)
         return params, opt_state, loss_value
 
@@ -189,7 +202,7 @@ def step(
     loss = []
     val_loss = []
 
-    opt_state = optimizer_provider.init(mc_initial_parameters)
+    opt_state = optim.init(mc_initial_parameters)
     parameters = mc_initial_parameters.copy()
 
     data_batch = data_batches(len_tr_idx, batch_size, batch_seed)

colibri/optax_optimizer.py

@@ -14,7 +14,7 @@
 
 
 def optimizer_provider(
-    optimizer="adam", optimizer_settings={}
+    optimizer="adam", optimizer_settings={}, clipnorm=6.073e-6
 ) -> optax._src.base.GradientTransformationExtraArgs:
     """
     Define the optimizer.
@@ -27,6 +27,9 @@ def optimizer_provider(
     optimizer_settings : dict, default = {}
         Dictionary containing the optimizer settings.
 
+    clipnorm : float, optional
+        If not None, gradients will be clipped to this norm.
+
     Returns
     -------
     optax._src.base.GradientTransformationExtraArgs
@@ -37,9 +40,13 @@ def optimizer_provider(
     if not "learning_rate" in optimizer_settings.keys():
         optimizer_settings["learning_rate"] = 5e-4
 
-    opt = getattr(optax, optimizer)
+    opt = getattr(optax, optimizer)(**optimizer_settings)
+
+    if clipnorm is not None:
+        log.info(f"Using clipnorm: {clipnorm}")
+        return optax.chain(optax.clip_by_global_norm(clipnorm), opt)
 
-    return opt(**optimizer_settings)
+    return opt
 
 
 def early_stopper(