Add function to reuse apriori-trained model

Author: SCiarella

benchmark.jl

@@ -11,6 +11,7 @@ end
 basedir = haskey(ENV, "DEEPDIP") ? ENV["DEEPDIP"] : @__DIR__
 outdir = joinpath(basedir, "output", "kolmogorov")
 confdir = joinpath(basedir, "configs/local")
+confdir = joinpath(basedir, "configs/snellius")
 @warn "Using configuration files from $confdir"
 compdir = joinpath(outdir, "comparison")
 ispath(compdir) || mkpath(compdir)

cnn_model_workflow.jl

@@ -176,6 +176,7 @@ end
 
 # Computational time
 docomp = conf["docomp"]
+docomp = false
 docomp && let
     comptime, datasize = 0.0, 0.0
     for seed in dns_seeds
@@ -241,6 +242,13 @@ closure_INS, θ_INS = NeuralClosure.cnn(;
 # Save parameters to disk after each run.
 # Plot training progress (for a validation data batch).
 
+# Check if it is asked to re-use the a-priori training from a different model
+if haskey(conf["priori"], "reuse")
+    reuse = conf["priori"]["reuse"]
+    @info "Reuse a-priori training from closure named: $reuse"
+    reusepriorfile(reuse, outdir, closure_name)
+end
+
 # Train
 for i = 1:ntrajectory
 	if i%numtasks == taskid -1

configs/snellius/cnn_1.yaml

@@ -10,7 +10,7 @@ params:
   tsim: 5.0
   savefreq: 50
   ndns: 4096
-  nles: [128]
+  nles: [64]
   filters: ["FaceAverage()"]
   icfunc: "(setup, psolver, rng) -> random_field(setup, T(0); kp=20, psolver, rng)"
   method: "RKMethods.Wray3(; T)"
@@ -24,7 +24,7 @@ seeds:
   prior: 345
   post: 456
 closure:
-  name: "cnn"
+  name: "cnn_project"
   type: cnn
   radii: [2, 2, 2, 2, 2]
   channels: [24, 24, 24, 24, 2]
@@ -42,9 +42,10 @@ posteriori:
   dotrain: true
   projectorders: "(ProjectOrder.Last, )"
   nepoch: 300
-  opt: "OptimiserChain(Adam(T(1.0e-3)), ClipGrad(0.1))"
+  opt: "OptimiserChain(Adam(T(1.0e-4)), ClipGrad(0.01))"
   nunroll: 5
   nunroll_valid: 10
+  nsamples: 1
   dt: 0.0001
   do_plot: false
   plot_train: false

configs/snellius/cnn_2.yaml

@@ -0,0 +1,52 @@
+docreatedata: true
+docomp: true
+ntrajectory: 8
+T: "Float32"
+params:
+  D: 2
+  lims: [0.0, 1.0]
+  Re: 6000.0
+  tburn: 0.5
+  tsim: 5.0
+  savefreq: 50
+  ndns: 4096
+  nles: [64]
+  filters: ["FaceAverage()"]
+  icfunc: "(setup, psolver, rng) -> random_field(setup, T(0); kp=20, psolver, rng)"
+  method: "RKMethods.Wray3(; T)"
+  bodyforce: "(dim, x, y, t) -> (dim == 1) * 5 * sinpi(8 * y)"
+  issteadybodyforce: true
+  processors: "(; log = timelogger(; nupdate=100))"
+  Δt: 0.00005
+seeds:
+  dns: 123456
+  θ_start: 234
+  prior: 345
+  post: 456
+closure:
+  name: "cnn_remove"
+  type: cnn
+  radii: [2, 2, 2, 2, 2]
+  channels: [24, 24, 24, 24, 2]
+  activations: ["tanh", "tanh", "tanh", "tanh", "identity"]
+  use_bias: [true, true, true, true, false]
+  rng: "Xoshiro(seeds.θ_start)"
+priori:
+  reuse: "cnn_project"
+  dotrain: true
+  nepoch: 10000
+  batchsize: 64
+  opt: "OptimiserChain(Adam(T(1.0e-3)), ClipGrad(0.1))"
+  do_plot: false
+  plot_train: false
+posteriori:
+  dotrain: true
+  projectorders: "(ProjectOrder.Last, )"
+  nepoch: 300
+  opt: "OptimiserChain(Adam(T(1.0e-4)), ClipGrad(0.01))"
+  nunroll: 5
+  nunroll_valid: 10
+  nsamples: 1
+  dt: 0.0001
+  do_plot: false
+  plot_train: false

configs/snellius/cnn_ins.yaml

@@ -0,0 +1,50 @@
+docreatedata: true
+docomp: true
+ntrajectory: 8
+T: "Float32"
+params:
+  D: 2
+  lims: [0.0, 1.0]
+  Re: 6000.0
+  tburn: 0.5
+  tsim: 5.0
+  savefreq: 50
+  ndns: 4096
+  nles: [64]
+  filters: ["FaceAverage()"]
+  icfunc: "(setup, psolver, rng) -> random_field(setup, T(0); kp=20, psolver, rng)"
+  method: "RKMethods.Wray3(; T)"
+  bodyforce: "(dim, x, y, t) -> (dim == 1) * 5 * sinpi(8 * y)"
+  issteadybodyforce: true
+  processors: "(; log = timelogger(; nupdate=100))"
+  Δt: 0.00005
+seeds:
+  dns: 123
+  θ_start: 234
+  prior: 345
+  post: 456
+closure:
+  name: "cnn_noproj"
+  type: cnn
+  radii: [2, 2, 2, 2, 2]
+  channels: [24, 24, 24, 24, 2]
+  activations: ["tanh", "tanh", "tanh", "tanh", "identity"]
+  use_bias: [true, true, true, true, false]
+  rng: "Xoshiro(seeds.θ_start)"
+priori:
+  dotrain: true
+  nepoch: 10000
+  batchsize: 64
+  opt: "OptimiserChain(Adam(T(1.0e-3)), ClipGrad(0.1))"
+  do_plot: false
+  plot_train: false
+posteriori:
+  dotrain: true
+  projectorders: "(ProjectOrder.Last, )"
+  nepoch: 300
+  opt: "OptimiserChain(Adam(T(1.0e-3)), ClipGrad(0.1))"
+  nunroll: 5
+  nunroll_valid: 10
+  dt: 0.0001
+  do_plot: false
+  plot_train: false

configs/snellius/conf_INS.yaml

@@ -12,7 +12,7 @@ params:
   tsim: 5.0
   savefreq: 50
   ndns: 4096
-  nles: [128]
+  nles: [64]
   filters: ["FaceAverage()"]
   icfunc: "(setup, psolver, rng) -> random_field(setup, T(0); kp=20, psolver, rng)"
   method: "RKMethods.Wray3(; T)"

src/Benchmark.jl

@@ -111,5 +111,6 @@ export _convert_to_single_index,
     plot_epost_vs_t
 
 export compute_eprior, compute_epost, compute_t_prior_inference
+export reusepriorfile
 
 end # module Benchmark

src/train.jl

@@ -26,7 +26,7 @@ function createdata(; params, seed, outdir, backend)
         @info "Creating data for" nles Φ seed
 
         data = NS.create_les_data_projected(
-            nchunks = 500;
+            nchunks = 8000;
             params...,
             rng = Xoshiro(seed),
             backend = backend,
@@ -44,6 +44,22 @@ function getpriorfile(outdir, closure_name, nles, filter)
     )
 end
 
+function reusepriorfile(reuse, outdir, closure_name)
+    reusepath = joinpath(outdir, "priortraining", reuse)
+    targetpath = joinpath(outdir, "priortraining", closure_name)
+    # If the reuse path exists, copy it to the target path
+    if ispath(reusepath)
+        @info "Reusing prior training from $(reusepath) to $(targetpath)"
+        ispath(targetpath) || mkpath(targetpath)
+        for file in readdir(reusepath, join = true)
+            @info "Copying prior training file $(file) to $(targetpath)"
+            cp(file, joinpath(targetpath, basename(file)); force = true)
+        end
+    else
+        @warn "Reuse path $(reusepath) does not exist. Not reusing prior training."
+    end
+end
+
 "Load a-priori training results from correct file names."
 loadprior(outdir, closure_name, nles, filters) = map(
     splat((nles, Φ) -> load_object(getpriorfile(outdir, closure_name, nles, Φ))),
@@ -104,7 +120,6 @@ function trainprior(;
         # Read the data in the format expected by the CoupledNODE
         data_train = load_data_set(outdir, nles, Φ, dns_seeds_train)
         data_valid = load_data_set(outdir, nles, Φ, dns_seeds_valid)
-
         @assert length(nles) == 1 "Only one nles for a-priori training"
         io_train = NS.create_io_arrays_priori(data_train, setup[1], device)
         io_valid = NS.create_io_arrays_priori(data_valid, setup[1], device)