MC dropout implementation, with scripts and slurms for reproduction

Author: rbonnetguerrini

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

@@ -195,8 +195,8 @@ def mc_dropout_predict(self, x=None, n_samples=100):
         (``training=True``), implementing MC Dropout inference.
 
         Each forward pass draws a fresh random binary mask on every Dropout
-        layer, so the spread across samples reflects the epistemic uncertainty
-        captured by the dropout regulariser.
+        layer, so the spread across samples reflects the uncertainty
+        captured by the dropout.
 
         Parameters
         ----------

n3fit/src/n3fit/mc_dropout_assemble.py

@@ -11,13 +11,13 @@
   --mode mean  (default)
     For every trained replica N write ONE exportgrid using the mean over all
     MC Dropout samples.  Gives 1 PDF member per trained replica.
-    Uncertainty comes from inter-replica diversity (same as standard NNPDF).
+    Uncertainty comes from replica diversity (same as standard NNPDF).
 
   --mode samples
     For every trained replica N write ONE exportgrid per MC Dropout sample.
     Gives n_samples PDF members per trained replica (sequential numbering).
     Uncertainty comes from the MC Dropout spread within a single trained
-    replica.  This is the mode you want to study MC Dropout uncertainty.
+    replica.  
     Example: 1 trained replica x 100 samples -> 100 PDF members.
 
 Usage

n3fit/src/n3fit/mc_dropout_inference.py

@@ -1,18 +1,18 @@
 """
-MC Dropout inference for NNPDF dropout fits - (fit on central value)
+MC Dropout inference for NNPDF dropout fits
 ======================================================================
 
 Loads the trained weights of one replica, runs N stochastic forward passes
 with dropout kept active (training=True), and computes the mean PDF over
-those passes.  The result is written as a numpy archive (.npz) containing:
+those passes.  The result contains:
 
   - ``x``           : x-grid, shape (n_x,)
   - ``mean``        : mean PDF,  shape (n_x, 14)
   - ``std``         : std  PDF,  shape (n_x, 14)
   - ``samples``     : all N samples, shape (N, n_x, 14)
   - ``flavours``    : LHAPDF PID list, shape (14,)
 
-Usage (from the project root, inside environment_nnpdf):
+Usage :
     python -m n3fit.mc_dropout_inference \\
         --fit-dir     nnpdf40-like-dropout-cluster \\
         --replica     1 \\
@@ -25,8 +25,7 @@
     nnfit/replica_<N>/weights.weights.h5   - saved Keras weights
     filter.yml or n3fit runcard            - architecture parameters
 
-The architecture is hardcoded from the runcard below and must match the
-saved weights.  If you change the runcard parameters, update the defaults.
+Warning: The architecture parameters are read from the n3fit runcard, not from the saved model.  
 """
 
 import argparse
@@ -63,10 +62,7 @@ def _get_xgrid():
     from n3fit.io.writer import XGRID  # noqa - re-raise if truly missing
     return XGRID
 
-
-# ---------------------------------------------------------------------------
-# ARCHITECTURE DEFAULTS : must match what was used during training
-# ---------------------------------------------------------------------------
+# Loading the architecture parameters from the fit runcard 
 DEFAULT_RUNCARD = (
     Path(__file__).resolve().parent.parent.parent  # n3fit/src/../..  : n3fit
     / "runcards" / "examples" / "nnpdf40-like-dropout-cluster.yml"
@@ -83,13 +79,13 @@ def _load_architecture(runcard_path):
     fitbasis = rc["fitting"]["fitbasis"]
 
     return dict(
-        nodes          = params["nodes_per_layer"],        # [25, 20, 8]
-        activations    = params["activation_per_layer"],   # ['tanh','tanh','linear']
-        initializer    = params["initializer"],            # 'glorot_normal'
-        architecture   = params["layer_type"],             # 'dense'
-        dropout_rate   = params.get("dropout", 0.0),      # 0.1
+        nodes          = params["nodes_per_layer"],        
+        activations    = params["activation_per_layer"],  
+        initializer    = params["initializer"],            
+        architecture   = params["layer_type"],             
+        dropout_rate   = params.get("dropout", 0.0),      
         flav_info      = basis,
-        fitbasis       = fitbasis,                         # 'EVOL'
+        fitbasis       = fitbasis,                         
     )