Merge pull request #2443 from NNPDF/load_weights_from_fit

Load weights from a fit

Author: scarlehoff

doc/sphinx/source/n3fit/runcard_detailed.rst

@@ -63,7 +63,7 @@ Monte Carlo replica data
 ------------------------
 The uncertainties on the PDFs are determined trough a Monte Carlo sampling procedure
 where each pseudodata replica is generated by adding some noise to the central value
-of the data points. 
+of the data points.
 In order to generate the pseudodata, the user can set the following parameter in the runcard:
 
 .. code-block:: yaml
@@ -74,8 +74,8 @@ The noise is generated by sampling from a normal distribution
 that is centred around zero. There are two options for the covariance matrix
 of the noise: the experimental covariance matrix or the t0 covariance matrix (default).
 
-In order to not use the t0-covariance matrix in the pseudodata generation, it is enough 
-to set 
+In order to not use the t0-covariance matrix in the pseudodata generation, it is enough
+to set
 
 .. code-block:: yaml
 
@@ -396,9 +396,12 @@ Save and load weights of the model
 
     save: "weights.h5"
     load: "weights.h5"
+    load_weights_from_fit: NNPDF40_nnlo_as_01180_qcd
 
 - ``save``: saves the weights of the PDF model in the selected file in the replica folder.
-- ``load``: loads the weights of the PDF model from the selected file.
+- ``load``: loads the weights of the PDF model from the selected file. Each replica will load the same weights from this file.
+- ``load_weights_from_fit``: loads the weights of the PDF model from a previous fit located in the `resultspath` directory.
+  This requires that the fit contains the saved weights. Each replica loads the weights from the corresponding replica of the specified fit.
 
 Since the weights depend only on the architecture of the Neural Network,
 it is possible to save the weights of a Neural Network trained with one set of hyperparameters and experiments

n3fit/src/n3fit/backends/keras_backend/MetaModel.py

@@ -395,6 +395,17 @@ def set_replica_weights(self, weights, i_replica=0):
             layer = self.get_layer(layer_type)
             set_layer_replica_weights(layer=layer, weights=weights[layer_type], i_replica=i_replica)
 
+    def set_replica_weights_from_file(self, model_file, i_replica=0):
+        """
+        Set the weights of replica i_replica from a model file.
+        Wrapper around ``set_replica_weights`` that creates a temporary single-replica-model
+        to get the weights in the appropiate format.
+        """
+        single_replica = self.single_replica_generator(i_replica)
+        single_replica.load_weights(model_file)
+        weights = single_replica.get_replica_weights(0)
+        self.set_replica_weights(weights, i_replica)
+
     def split_replicas(self):
         """
         Split the single multi-replica model into a list of separate single replica models,

n3fit/src/n3fit/checks.py

@@ -16,17 +16,17 @@
 
 NN_PARAMETERS = ["nodes_per_layer", "optimizer", "activation_per_layer"]
 HYPEROPTIMIZED_PARAMETERS = [
-    "nodes_per_layer" , 
-    "optimizer", 
-    "learning_rate", 
+    "nodes_per_layer",
+    "optimizer",
+    "learning_rate",
     "clipnorm",
     "activation_per_layer",
     "initializer",
     "epochs",
     "stopping_patience",
     "layer_type",
-    "dropout"
-    ]
+    "dropout",
+]
 
 
 def _is_floatable(num):
@@ -40,13 +40,17 @@ def _is_floatable(num):
     except (ValueError, TypeError):
         return False
 
+
 def check_hyperopt_parameters(parameters, trial_specs):
     tmp = []
     for param in HYPEROPTIMIZED_PARAMETERS:
         if param in parameters.keys():
             tmp.append(param)
     if tmp:
-        raise CheckError(f"Parameters {tmp} already contained in the hyperoptimization scan. Please remove them from the parameters namespace.")        
+        raise CheckError(
+            f"Parameters {tmp} already contained in the hyperoptimization scan. Please remove them from the parameters namespace."
+        )
+
 
 # Checks on the NN parameters
 def check_existing_parameters(parameters):
@@ -228,11 +232,38 @@ def check_model_file(save, load):
             raise CheckError(f"Model file {load} seems to be empty")
 
 
+def check_load_fits_from_weight_file(load_weights_from_fit, load, load_weights_dict, replicas):
+    """Checks whether the load_weights_from_fit option is correctly defined.
+    And whether the requested replica can be loaded
+    """
+    if load_weights_from_fit is not None:
+        if load is not None:
+            raise CheckError(
+                "Cannot use both `load` and `load_weights_from_fit` options at the same time, please select only one of them"
+            )
+    if replicas is not None and load_weights_from_fit is not None:
+        missing_replicas = set(replicas) - set(load_weights_dict.keys())
+        if missing_replicas:
+            raise CheckError(
+                f"Not all replicas requested have weights to be loaded, missing: {missing_replicas}"
+            )
+
+
 @make_argcheck
-def wrapper_check_NN(tensorboard, save, load, parameters, trial_specs):
+def wrapper_check_NN(
+    tensorboard,
+    save,
+    load,
+    load_weights_from_fit,
+    load_weights_dict,
+    parameters,
+    trial_specs,
+    replicas,
+):
     """Wrapper function for all NN-related checks"""
     check_tensorboard(tensorboard)
     check_model_file(save, load)
+    check_load_fits_from_weight_file(load_weights_from_fit, load, load_weights_dict, replicas)
     check_lagrange_multipliers(parameters, "integrability")
     check_lagrange_multipliers(parameters, "positivity")
     if trial_specs:
@@ -489,7 +520,7 @@ def check_deprecated_options(fitting, parameters):
         raise CheckError(
             "`interpolation_points` no longer accepted, please change to `feature_scaling_points`"
         )
-        
+
 
 @make_argcheck
 def check_photonQED_exists(theoryid, fiatlux):

n3fit/src/n3fit/model_trainer.py

@@ -27,10 +27,11 @@
 from n3fit.stopping import Stopping
 from n3fit.vpinterface import N3PDF, compute_hyperopt_metrics
 from validphys.core import DataGroupSpec
+from validphys.loader import Loader
 from validphys.photon.compute import Photon
 
 log = logging.getLogger(__name__)
-
+l = Loader()
 # Threshold defaults
 # Any partition with a chi2 over the threshold will discard its hyperparameters
 HYPER_THRESHOLD = 50.0
@@ -113,6 +114,7 @@ def __init__(
         lux_params=None,
         replicas=None,
         trials=None,
+        load_weights_dict=None,
     ):
         """
         Parameters
@@ -179,6 +181,7 @@ def __init__(
         else:
             self.max_cores = max_cores
         self.model_file = model_file
+        self.load_weights_dict = load_weights_dict
         self.print_summary = True
         self.mode_hyperopt = False
         self.impose_sumrule = sum_rules
@@ -987,12 +990,20 @@ def hyperparametrizable(self, params):
             # together with pdf model generated above
             models = self._model_generation(xinput, pdf_model, partition, k)
 
-            # Only after model generation, apply possible weight file
-            # Starting every replica with the same weights
+            # After model generation, apply possible weights files.
+            # The possibilities are a single model file (`load:`)
+            # so that every replica starts with the same weights
+            # or a weight per replica from `load_weights_from_fit`
             if self.model_file:
                 log.info("Applying model file %s", self.model_file)
                 pdf_model.load_identical_replicas(self.model_file)
 
+            if self.load_weights_dict:
+                for replica in self.replicas:
+                    weights_path = self.load_weights_dict[replica]
+                    log.info("Loading weights from path: " + str(weights_path))
+                    pdf_model.set_replica_weights_from_file(weights_path)
+
             if k > 0:
                 # Reset the positivity and integrability multipliers
                 pos_and_int = self.training["posdatasets"] + self.training["integdatasets"]

n3fit/src/n3fit/n3fit_checks_provider.py

@@ -23,13 +23,16 @@ def n3fit_checks_action(
     parameters,
     save=None,
     load=None,
+    load_weights_from_fit=None,
+    load_weights_dict,
     hyperscan_config=None,
     hyperopt=None,
     kfold=None,
     tensorboard=None,
     parallel_models=True,
     double_precision=False,
     trial_specs=None,
+    replicas=None,
 ):
     return
 

n3fit/src/n3fit/performfit.py

@@ -28,6 +28,7 @@ def performfit(
     basis,
     fitbasis,
     positivity_bound,
+    load_weights_dict,
     sum_rules=True,
     parameters,
     replica_path,
@@ -201,6 +202,7 @@ def performfit(
             lux_params=fiatlux,
             replicas=replica_idxs,
             trials=trials,
+            load_weights_dict=load_weights_dict
         )
 
         # This is just to give a descriptive name to the fit function

validphys2/src/validphys/config.py

@@ -165,6 +165,29 @@ def parse_pdf(self, name, unpolarized_bc=None):
 
         return pdf
 
+    def produce_load_weights_dict(self, load_weights_from_fit=None):
+        if load_weights_from_fit is None:
+            return None
+        try:
+            fit_object = self.loader.check_fit(load_weights_from_fit)
+            fit_folder = fit_object.path / "nnfit"
+            weights_name = fit_object.as_input().get("save")
+            if weights_name is None:
+                raise LoadFailedError(f"{load_weights_from_fit} does not have saved weights")
+            # Correct for extension already included
+            if weights_name.endswith(".h5"):
+                weights_name = weights_name[:-3]
+            weights_dict = {}
+            for p in fit_folder.glob(f"replica_*/{weights_name}.weights.h5"):
+                replica_folder = p.parent.name
+                replica_index = int(replica_folder.split("_")[1])
+                weights_dict[replica_index] = p
+            if not weights_dict:
+                raise LoadFailedError(f"No weights found for: {load_weights_from_fit}")
+            return weights_dict
+        except LoadFailedError as e:
+            raise ConfigError(str(e), load_weights_from_fit, self.loader.available_fits) from e
+
     @element_of("unpolarized_bcs")
     @_id_with_label
     def parse_unpolarized_bc(self, name):