Save best parameters

Author: achiefa

n3fit/src/n3fit/backends/keras_backend/callbacks.py

@@ -212,16 +212,25 @@ class StoreCallback(CallbackStep):
             Save every this many epochs (default: 100)
     """
 
-    def __init__(self, pdf_model, replica_paths, check_freq=100):
+    def __init__(self, pdf_model, replica_paths, stopping_object, check_freq=100):
         super().__init__()
         self.check_freq = check_freq
         self.pdf_model = pdf_model
         self.weight_dirs = []
+        self.stopping_object = stopping_object
         for path in replica_paths:
             weight_dir = path / "parameters"
             weight_dir.mkdir(parents=True, exist_ok=True)
             self.weight_dirs.append(weight_dir)
 
+    def _save_weights(self, epoch, tr_weights, weight_dir):
+
+        filepath = weight_dir / f"params_{epoch+1}.npz"
+        # save parameters as expected by colibri
+        trainable_weights_flat = np.concatenate([np.asarray(w).flatten() for w in tr_weights])
+        np.savez(filepath, params=trainable_weights_flat)
+        log.info(f"Saved parameters at epoch {epoch+1} in {filepath}")
+
     def on_step_end(self, epoch, logs=None):
         """Function to be called at the end of every epoch
         Every ``check_freq`` number of epochs, the parameters of the model will
@@ -230,14 +239,15 @@ def on_step_end(self, epoch, logs=None):
         if ((epoch + 1) % self.check_freq) == 0:
             pdf_replicas = self.pdf_model.split_replicas()
             for replica_model, weight_dir in zip(pdf_replicas, self.weight_dirs):
-                filepath = weight_dir / f"params_{epoch+1}.npz"
-                # save parameters as expected by colibri
-                trainable_weights_flat = np.concatenate(
-                    [w.numpy().flatten() for w in replica_model.trainable_weights]
-                )
-                np.savez(filepath, params=trainable_weights_flat)
-                # replica_model.save_weights(filepath)
-                log.info(f"Saved parameters at epoch {epoch+1} in {filepath}")
+                weights = replica_model.trainable_weights
+                self._save_weights(epoch, weights, weight_dir)
+
+    def on_train_end(self, logs=None):
+        """Store the best parameters"""
+        for idx, weight_dir in enumerate(self.weight_dirs):
+            best_epoch = self.stopping_object._best_epochs[idx]
+            best_weights = self.stopping_object._best_weights[idx]['all_NNs']
+            self._save_weights(best_epoch, best_weights, weight_dir)
 
 
 def gen_tensorboard_callback(log_dir, profiling=False, histogram_freq=0):

n3fit/src/n3fit/model_trainer.py

@@ -752,7 +752,10 @@ def _train_and_fit(self, training_model, stopping_object, epochs=100) -> bool:
                 self.replica_path.parent / f"fit_replicas/replica_{r}" for r in self.replicas
             ]
             checpoint_callback = callbacks.StoreCallback(
-                pdf_model=pdf_model, replica_paths=replica_paths, check_freq=self.checkpoint_freq
+                pdf_model=pdf_model,
+                replica_paths=replica_paths,
+                check_freq=self.checkpoint_freq,
+                stopping_object=stopping_object,
             )
             callback_list.append(checpoint_callback)