Add regularization term

Author: SCiarella

configs/bck_snell/att_1.yaml

@@ -25,7 +25,7 @@ seeds:
   prior: 345
   post: 456
 closure:
-  name: "attention"
+  name: "attention_reg"
   type: attentioncnn
   radii: [2, 2, 2, 2, 2]
   channels: [24, 24, 24, 24, 2]
@@ -42,19 +42,21 @@ closure:
   rng: "Xoshiro(seeds.θ_start)"
 priori:
   dotrain: true
-  nepoch: 50000
+  nepoch: 5000
   batchsize: 64
   opt: "OptimiserChain(Adam(T(1.0e-3)), ClipGrad(0.1))"
   do_plot: false
   plot_train: false
+  lambda: 0.001
 posteriori:
   dotrain: true
   projectorders: "(ProjectOrder.Last, )"
   nepoch: 1500
-  opt: "OptimizationCMAEvolutionStrategy.CMAEvolutionStrategyOpt()"
+  opt: "OptimiserChain(Adam(T(1.0e-4)), ClipGrad(0.01))"
   nunroll: 5
   nunroll_valid: 10
   dt: 0.0001
   nsamples: 1
   do_plot: false
   plot_train: false
+  lambda: 0.001

configs/bck_snell/cnn_noproj.yaml

@@ -25,7 +25,7 @@ seeds:
   prior: 345
   post: 456
 closure:
-  name: "cnn_noproj"
+  name: "NoProjection"
   type: cnn
   radii: [2, 2, 2, 2, 2]
   channels: [24, 24, 24, 24, 2]

extra_model_workflow.jl

@@ -299,6 +299,7 @@ let
         plot_train = conf["priori"]["plot_train"],
         nepoch,
         dataproj = conf["dataproj"],
+        λ = conf["priori"]["lambda"],
     )
 end
 end
@@ -411,6 +412,7 @@ let
         sensealg = sensealg,
         sciml_solver = sciml_solver,
         dataproj = conf["dataproj"],
+        λ = conf["posteriori"]["lambda"],
     )
 end
 end

src/train.jl

@@ -110,7 +110,8 @@ function trainprior(;
     do_plot = false,
     plot_train = false,
     nepoch,
-    dataproj
+    dataproj,
+    λ = nothing,
 )
     device(x) = adapt(params.backend, x)
     itotal = 0
@@ -159,8 +160,9 @@ function trainprior(;
         )
         train_data_priori = dataloader_prior()
 
-        loss_priori_lux(closure, θ, st, train_data_priori)
-        loss = loss_priori_lux
+        # Trigger the loss once and wrap it for the expected Lux interface
+        loss_priori_lux(closure, θ, st, train_data_priori, λ)
+        loss(model, param, state, data) = loss_priori_lux(model, param, state, data, λ)
 
         if loadcheckpoint && isfile(checkfile)
             callbackstate, trainstate, epochs_trained =
@@ -272,6 +274,7 @@ function trainpost(;
     sensealg = nothing,
     sciml_solver = nothing,
     dataproj,
+    λ = nothing,
 )
     device(x) = adapt(params.backend, x)
     itotal = 0
@@ -325,6 +328,7 @@ function trainpost(;
             griddims,
             inside,
             dt;
+            λ = λ,
             ensemble = nsamples > 1,
             sciml_solver = sciml_solver,
             sensealg = sensealg,