add mfvi as vi baseline

Author: zuhengxu

example/real_data_expt/baseline/run_mfvi.jl

@@ -0,0 +1,89 @@
+using Random, Distributions
+using LinearAlgebra
+using LogDensityProblems, LogDensityProblemsAD
+using JLD2
+
+using MixFlow 
+const MF = MixFlow
+
+
+include(joinpath(@__DIR__, "../../Model.jl"))
+include(joinpath(@__DIR__, "../../julia_env/flowlayer.jl"))
+
+function run_baseline(
+    seed, name::String, lr; 
+    batchsize::Int = 64, niters::Int= 50_000, show_progress=true,
+    nsample_eval::Int=1024, save_jld::Bool=true,
+)
+    Random.seed!(seed)
+
+    @info "load model $(name)"
+    target, dims, ad = load_model(name)
+
+    @info "learning mfvi for $(name), dims = $(dims)"
+    dim = LogDensityProblems.dimension(target)
+    logp = Base.Fix1(LogDensityProblems.logdensity, target)
+
+    q₀ = MvNormal(zeros(dim), I)
+    flow =
+        Bijectors.transformed(q₀, Bijectors.Shift(zeros(dim)) ∘ Bijectors.Scale(ones(dim)))
+    
+    cb(iter, opt_stats, re, θ) = (sample_per_iter = sample_per_iter, ad = ad)
+    checkconv(iter, stat, re, θ, st) = _is_nan_or_inf(stat.loss) || (stat.gradient_norm < 1e-3)
+
+    time = @elapsed begin
+    flow_trained, stats, _ = train_flow(
+        NormalizingFlows.elbo,
+        flow,
+        logp,
+        batchsize;
+        max_iters=niters,
+        optimiser=Optimisers.Adam(lr),
+        ADbackend=ad,
+        show_progress=show_progress,
+        hasconverged=checkconv,
+        callback=cb,
+    )
+    end
+    @info "Training finished"
+
+    # if early stop due to NaN or Inf, return NaN for all
+    if _is_nan_or_inf(stats[end].loss)
+        println("Training failed: loss is NaN or Inf")
+        return DataFrame(
+            time = NaN,
+            elbo = NaN,
+            logZ = NaN,
+            ess = NaN,
+        )
+    end
+
+    # losses = map(x -> x.loss, stats)
+    # try and if error happens, return NaN
+    el, logz, es = flow_sample_eval(logp, flow_trained; nsample = nsample_eval)
+    
+    # save the trained flow
+    if save_jld
+        res_dir = joinpath(@__DIR__, "result/")
+
+        if !isdir(res_dir)
+            mkdir(res_dir)
+        end
+
+        JLD2.save(
+            joinpath(res_dir, "$(name)_mfvi_$(lr)_$(seed).jld2"),
+            "flow", flow_trained,
+            "batchsize", batchsize,
+            "seed", seed,
+        )
+    end
+    
+    df = DataFrame(
+        time = time,
+        elbo=el,
+        logZ=logz,
+        ess=es,
+    )
+
+    return df
+end

example/real_data_expt/baseline/run_mfvi.nf

@@ -0,0 +1,87 @@
+include { crossProduct; filed; deliverables } from '../../nf-nest/cross.nf'
+include { instantiate; precompile; activate } from '../../nf-nest/pkg.nf'
+include { combine_csvs; } from '../../nf-nest/combine.nf'
+
+params.dryRun = false
+params.n_sample_eval = params.dryRun ? 8 : 1024
+params.nrunThreads = 1
+
+def julia_env = file("${moduleDir}/../../julia_env")
+def julia_script = file("${moduleDir}/run_ais.jl")
+
+def variables = [
+    target: ["Sonar", "Brownian", "TReg", "SparseRegression" ,"LGCP"],
+    lr: ["1e-3"],
+    batchsize: [64],
+    niters: [50000],
+    seed: 1..10,
+]
+
+workflow {
+    compiled_env = instantiate(julia_env) | precompile
+    configs = crossProduct(variables, params.dryRun)
+    combined = run_simulation(compiled_env, configs) | combine_csvs
+    // plot(compiled_env, plot_script, combined)
+   final_deliverable(compiled_env, combined)
+}
+
+
+process run_simulation {
+    debug false 
+    memory { 30.GB * Math.pow(2, task.attempt-1) }
+    time { 24.hour * Math.pow(2, task.attempt-1) } 
+    cpus 1 
+    errorStrategy { task.attempt < 2 ? 'retry' : 'ignore' } 
+    input:
+        path julia_env 
+        val config 
+    output:
+        path "${filed(config)}"
+    """
+    ${activate(julia_env,params.nrunThreads)}
+
+    include("$julia_script")
+
+    # get configurations
+    seed = ${config.seed}
+    name = "${config.target}"
+    niters = ${config.niters}
+    bs = ${config.batchsize} 
+    lr = ${config.lr}
+
+    # run simulation
+    try
+        df = run_baseline(
+            seed, name, lr; 
+            batchsize=bs, niters=niters, show_progress=false,
+            nsample_eval=${params.n_sample_eval},
+            save_jld = true
+        )
+    catch e
+        df = DataFrame(
+            time = NaN,
+            elbo = NaN,
+            logZ = NaN,
+            ess = NaN,
+        )
+    end
+
+    
+    # store output
+    mkdir("${filed(config)}")
+    CSV.write("${filed(config)}/summary.csv", df)
+    """
+}
+
+
+process final_deliverable {
+    input:
+        path julia_env 
+        path combined_csvs_folder 
+    output:
+        path combined_csvs_folder
+    publishDir "${deliverables(workflow, params)}", mode: 'copy', overwrite: true
+    """
+    ${activate(julia_env)}
+    """
+}