model tuning

experiments/Manifest.toml

@@ -805,9 +805,9 @@ uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820"
 
 [[deps.GLM]]
 deps = ["Distributions", "LinearAlgebra", "Printf", "Reexport", "SparseArrays", "SpecialFunctions", "Statistics", "StatsAPI", "StatsBase", "StatsFuns", "StatsModels"]
-git-tree-sha1 = "97829cfda0df99ddaeaafb5b370d6cab87b7013e"
+git-tree-sha1 = "273bd1cd30768a2fddfa3fd63bbc746ed7249e5f"
 uuid = "38e38edf-8417-5370-95a0-9cbb8c7f171a"
-version = "1.8.3"
+version = "1.9.0"
 
 [[deps.GPUArrays]]
 deps = ["Adapt", "GPUArraysCore", "LLVM", "LinearAlgebra", "Printf", "Random", "Reexport", "Serialization", "Statistics"]
@@ -2389,9 +2389,9 @@ uuid = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 version = "0.3.5"
 
 [[deps.TupleTools]]
-git-tree-sha1 = "c8cdc29448afa1a306419f5d1c7af0854c171c80"
+git-tree-sha1 = "155515ed4c4236db30049ac1495e2969cc06be9d"
 uuid = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
-version = "1.4.1"
+version = "1.4.3"
 
 [[deps.URIs]]
 git-tree-sha1 = "b7a5e99f24892b6824a954199a45e9ffcc1c70f0"

experiments/circles.jl

@@ -3,16 +3,11 @@ dataname = "Circles"
 n_obs = Int(1000 / (1.0 - TEST_SIZE))
 counterfactual_data, test_data = train_test_split(load_circles(n_obs; noise=0.05, factor=0.5); test_size=TEST_SIZE)
 
+# Model tuning:
+model_tuning_params = DEFAULT_MODEL_TUNING_SMALL
+
 # Tuning parameters:
-tuning_params = (
-    nsamples=[10, 50, 100],
-    niter_eccco=[20, 50, 100],
-    Λ=[
-        [0.1, 0.1, 0.1],
-        [0.1, 0.2, 0.2],
-        [0.1, 0.5, 0.5],
-    ]
-)
+tuning_params = DEFAULT_GENERATOR_TUNING
 
 # Parameter choices:
 params = (

experiments/experiment.jl

@@ -38,6 +38,7 @@ Base.@kwdef struct Experiment
     train_parallel::Bool = false
     reg_strength::Real = 0.1
     niter_eccco::Union{Nothing,Int} = nothing
+    model_tuning_params::Tuple = DEFAULT_MODEL_TUNING_SMALL
 end
 
 "A container to hold the results of an experiment."
@@ -88,13 +89,18 @@ function run_experiment(exper::Experiment; save_output::Bool=true, only_models::
     end
     outcome = ExperimentOutcome(exper, nothing, nothing, nothing)
 
-    # Models
-    train_models!(outcome, exper)
+    # Model tuning:
+    if TUNE_MODEL
+        mach = tune_model(exper)
+        return mach
+    end
 
+    # Model training:
+    train_models!(outcome, exper)
     # Return if only models are needed:
     !only_models || return outcome
 
-    # Benchmark
+    # Benchmark:
     benchmark!(outcome, exper)
     if is_multi_processed(exper)
         MPI.Barrier(exper.parallelizer.comm)

experiments/linearly_separable.jl

@@ -6,16 +6,11 @@ counterfactual_data, test_data = train_test_split(
     test_size=TEST_SIZE
 )
 
+# Model tuning:
+model_tuning_params = DEFAULT_MODEL_TUNING_SMALL
+
 # Tuning parameters:
-tuning_params = (
-    nsamples=[10, 50, 100],
-    niter_eccco=[20, 50, 100],
-    Λ=[
-        [0.1, 0.1, 0.1],
-        [0.1, 0.2, 0.2],
-        [0.1, 0.5, 0.5],
-    ]
-)
+tuning_params = DEFAULT_GENERATOR_TUNING
 
 # Parameter choices:
 params = (

experiments/model_tuning.jl

-"An MLP builder that is more easily tunable."
-mutable struct TuningBuilder <: MLJFlux.Builder
-    n_hidden::Int
-    n_layers::Int
-end
-
-"Outer constructor."
-TuningBuilder(;n_hidden=32, n_layers=3) = TuningBuilder(n_hidden, n_layers)
+"""
+    tune_model(exper::Experiment; kwargs...)
 
-function MLJFlux.build(nn::TuningBuilder, rng, n_in, n_out)
-    hidden = ntuple(i -> nn.n_hidden, nn.n_layers)
-    return MLJFlux.build(MLJFlux.MLP(hidden=hidden), rng, n_in, n_out)
+Tunes MLP in place and saves the tuned model to disk.
+"""
+function tune_model(exper::Experiment; kwargs...)
+    if !(is_multi_processed(exper) && MPI.Comm_rank(exper.parallelizer.comm) != 0)
+        @info "Tuning models."
+        # Output path:
+        model_tuning_path = mkpath(joinpath(DEFAULT_OUTPUT_PATH, "tuned_model"))
+        # Simple MLP:
+        mod = NeuralNetworkClassifier(
+            builder=default_builder(),
+            epochs=exper.epochs,
+            batch_size=batch_size(exper),
+            finaliser=exper.finaliser,
+            loss=exper.loss,
+            acceleration=CUDALibs(),
+        )
+        mach = tune_model(mod, X, y; tuning_params=exper.model_tuning_params, measure=exper.model_measures, kwargs...)
+        Serialization.serialize(joinpath(model_tuning_path, "$(exper.save_name).jls"), mach)
+    end
+    return mach
 end
 
 """
@@ -17,7 +28,7 @@ end
 
 Tunes a model by performing a grid search over the parameters specified in `tuning_params`.
 """
-function tune_model(mod::Supervised, X, y; tuning_params::NamedTuple, kwargs...)
+function tune_model(mod::Supervised, X, y; tuning_params::NamedTuple, measure=MODEL_MEASURES, kwargs...)
 
     ranges = []
 
@@ -36,6 +47,7 @@ function tune_model(mod::Supervised, X, y; tuning_params::NamedTuple, kwargs...)
     self_tuning_mod = TunedModel(
         model=mod,
         range=ranges,
+        measure=measure,
         kwargs...
     )
 
@@ -44,4 +56,5 @@ function tune_model(mod::Supervised, X, y; tuning_params::NamedTuple, kwargs...)
 
     return mach
 
-end
\ No newline at end of file
+end
+

experiments/models/default_models.jl

+"An MLP builder that is more easily tunable."
+mutable struct TuningBuilder <: MLJFlux.Builder
+    n_hidden::Int
+    n_layers::Int
+    activation::Function
+end
+
+"Outer constructor."
+TuningBuilder(; n_hidden=32, n_layers=3, activation=Flux.swish) = TuningBuilder(n_hidden, n_layers, activation)
+
+function MLJFlux.build(nn::TuningBuilder, rng, n_in, n_out)
+    hidden = ntuple(i -> nn.n_hidden, nn.n_layers)
+    return MLJFlux.build(MLJFlux.MLP(hidden=hidden, σ=nn.activation), rng, n_in, n_out)
+end
+
 """
     default_builder(n_hidden::Int=16, activation::Function=Flux.swish)
 
 Default builder for MLPs.
 """
-function default_builder(n_hidden::Int=16, activation::Function=Flux.swish)
-    builder = MLJFlux.MLP(
-        hidden=(n_hidden, n_hidden, n_hidden),
-        σ=activation
-    )
+function default_builder(n_hidden::Int=16, n_layers::Int=3, activation::Function=Flux.swish)
+    builder = TuningBuilder(n_hidden=n_hidden, n_layers=n_layers, activation=activation)
     return builder
 end
 

experiments/moons.jl

@@ -3,16 +3,11 @@ dataname = "Moons"
 n_obs = Int(2500 / (1.0 - TEST_SIZE))
 counterfactual_data, test_data = train_test_split(load_moons(n_obs); test_size=TEST_SIZE)
 
+# Model tuning:
+model_tuning_params = DEFAULT_MODEL_TUNING_SMALL
+
 # Tuning parameters:
-tuning_params = (
-    nsamples=[10, 50, 100],
-    niter_eccco=[20, 50, 100],
-    Λ=[
-        [0.1, 0.1, 0.1],
-        [0.1, 0.2, 0.2],
-        [0.1, 0.5, 0.5],
-    ]
-)
+tuning_params = DEFAULT_GENERATOR_TUNING
 
 # Parameter choices:
 params = (

experiments/setup_env.jl

@@ -141,4 +141,33 @@ const N_IND = n_individuals
 const N_IND_SPECIFIED = n_ind_specified
 
 "Boolean flag to check if grid search was specified."
-const GRID_SEARCH = "grid_search" ∈ ARGS
\ No newline at end of file
+const GRID_SEARCH = "grid_search" ∈ ARGS
+
+"Generator tuning parameters."
+const DEFAULT_GENERATOR_TUNING = (
+    nsamples=[10, 100],
+    niter_eccco=[10, 100],
+    Λ=[
+        [0.1, 0.1, 0.1],
+        [0.1, 0.2, 0.2],
+        [0.1, 0.5, 0.5],
+    ],
+    reg_strength=[0.0, 0.1, 0.5],
+)
+
+"Boolean flag to check if model tuning was specified."
+const TUNE_MODEL = "tune_model" ∈ ARGS
+
+"Model tuning parameters for small datasets."
+const DEFAULT_MODEL_TUNING_SMALL = (
+    n_hidden=[16, 32, 64],
+    n_layers=[1, 2, 3],
+    activation=[Flux.relu, Flux.swish],
+)
+
+"Model tuning parameters for large datasets."
+const DEFAULT_MODEL_TUNING_LARGE = (
+    n_hidden=[32, 64, 128],
+    n_layers=[2, 3, 5],
+    activation=[Flux.relu, Flux.swish],
+)