|
| 1 | +# Hyperparameter Optimization |
| 2 | + |
| 3 | +GraphNetSim integrates with [Optuna.jl](https://github.com/una-auxme/Optuna.jl) to automate the search |
| 4 | +for good GNN-simulator hyperparameters. This page documents the workflow implemented by the runnable |
| 5 | +example |
| 6 | +[`example/DamBreakSmall/DamBreakSmall_optuna.jl`](https://github.com/una-auxme/GraphNetSim.jl/tree/main/example/DamBreakSmall/DamBreakSmall_optuna.jl); |
| 7 | +the same pattern applies to any dataset. |
| 8 | + |
| 9 | +## Requirements |
| 10 | + |
| 11 | +Optuna is an extra dependency, so activate an environment that provides it. The DamBreakSmall example |
| 12 | +ships its own `Project.toml` with Optuna, OrdinaryDiffEq, and Optimisers, so run the script with that |
| 13 | +environment from the repository root: |
| 14 | + |
| 15 | +```bash |
| 16 | +julia --project=example/DamBreakSmall example/DamBreakSmall/DamBreakSmall_optuna.jl |
| 17 | +``` |
| 18 | + |
| 19 | +## Workflow overview |
| 20 | + |
| 21 | +The workflow has four parts: a **persistent study**, an **objective** that trains one model per trial, |
| 22 | +an **ask/tell loop** that samples the search space, and **result inspection**. A single trial trains a |
| 23 | +GNN with sampled hyperparameters and reports the best validation loss returned by |
| 24 | +[`train_network`](@ref GraphNetSim.train_network) — that is, the loss of the periodic ODE rollout on |
| 25 | +the validation split. Minimizing that value across trials is the optimization objective. |
| 26 | + |
| 27 | +### 1. A persistent, resumable study |
| 28 | + |
| 29 | +The study is backed by a SQLite database and an on-disk artifact store, so re-running the script |
| 30 | +continues an existing study rather than starting over: |
| 31 | + |
| 32 | +```julia |
| 33 | +storage_url = create_sqlite_url(database_url, database_name) |
| 34 | +storage = RDBStorage(storage_url) |
| 35 | +artifact_store = FileSystemArtifactStore(artifact_path) |
| 36 | + |
| 37 | +study = Study( |
| 38 | + study_name, |
| 39 | + artifact_store, |
| 40 | + storage; |
| 41 | + sampler=TPESampler(), # Tree-structured Parzen Estimator |
| 42 | + pruner=MedianPruner(5, 1), # stop trials worse than the running median |
| 43 | + direction="minimize", |
| 44 | + load_if_exists=true, # resume an existing study of the same name |
| 45 | +) |
| 46 | +``` |
| 47 | + |
| 48 | +- **Sampler** — `TPESampler` models the relationship between hyperparameters and loss and proposes |
| 49 | + promising configurations; swap in another sampler to change the search strategy. |
| 50 | +- **Pruner** — `MedianPruner` terminates unpromising trials early (after a startup grace period) by |
| 51 | + comparing a trial's reported loss against previous trials. |
| 52 | +- **`load_if_exists=true`** — combined with the SQLite storage, this is what makes the run resumable. |
| 53 | + |
| 54 | +### 2. The objective — one training run per trial |
| 55 | + |
| 56 | +The objective converts sampled parameters into a training configuration, runs |
| 57 | +[`train_network`](@ref GraphNetSim.train_network), and returns the best validation loss. Each trial |
| 58 | +trains into a fresh temporary checkpoint directory so trials do not interfere: |
| 59 | + |
| 60 | +```julia |
| 61 | +function objective(trial::Trial; params) |
| 62 | + cp_path = mktempdir() |
| 63 | + |
| 64 | + opt = if params[:optimizer] == "Adam" |
| 65 | + Adam(params[:lr]) |
| 66 | + elseif params[:optimizer] == "AdamW" |
| 67 | + AdamW(; eta=params[:lr], lambda=params[:weight_decay]) |
| 68 | + else |
| 69 | + RAdam(params[:lr]) |
| 70 | + end |
| 71 | + |
| 72 | + min_val_loss = train_network( |
| 73 | + opt, ds_path, cp_path; |
| 74 | + training_strategy=DerivativeTraining(; |
| 75 | + random=params[:random_sampling], window_size=params[:window_size] |
| 76 | + ), |
| 77 | + steps=n_steps, checkpoint=cp_interval, |
| 78 | + mps=params[:mps], layer_size=params[:layer_size], hidden_layers=params[:hidden_layers], |
| 79 | + noise_stddevs=[params[:noise_std]], |
| 80 | + norm_steps=0, norm_type=params[:norm_type], |
| 81 | + optimizer_learning_rate_start=params[:lr], |
| 82 | + optimizer_learning_rate_stop=params[:lr] * params[:lr_decay_ratio], |
| 83 | + # ... fixed args: types_updated, types_noisy, solver_valid, use_cuda, ... |
| 84 | + ) |
| 85 | + |
| 86 | + report(trial, Float64(min_val_loss), 1) # feed the pruner |
| 87 | + should_prune(trial) && return nothing |
| 88 | + |
| 89 | + upload_artifact(study, trial, Dict(String(k) => v for (k, v) in pairs(params))) |
| 90 | + return Float64(min_val_loss) |
| 91 | +end |
| 92 | +``` |
| 93 | + |
| 94 | +`report` hands the trial's loss to the pruner; `should_prune` then decides whether to abandon it; |
| 95 | +`upload_artifact` records the trial's hyperparameters for later inspection. |
| 96 | + |
| 97 | +### 3. The ask/tell loop — sampling the search space |
| 98 | + |
| 99 | +Each iteration `ask`s the study for a trial, draws hyperparameters with the `suggest_*` family, runs |
| 100 | +the objective, and `tell`s the study the result (or that it was pruned): |
| 101 | + |
| 102 | +```julia |
| 103 | +trial = ask(study) |
| 104 | + |
| 105 | +mps = suggest_int(trial, "mps", 3, 10) |
| 106 | +layer_size = suggest_categorical(trial, "layer_size", [32, 64, 128]) |
| 107 | +lr = suggest_float(trial, "lr", 1.0e-5, 1.0e-3; log=true) # log-scale |
| 108 | +norm_type = Symbol(suggest_categorical(trial, "norm_type", ["minmax", "meanstd"])) |
| 109 | +# ... remaining suggestions ... |
| 110 | + |
| 111 | +params = (; mps, layer_size, lr=Float32(lr), norm_type, #= ... =#) |
| 112 | +score = objective(trial; params) |
| 113 | + |
| 114 | +if isnothing(score) |
| 115 | + tell(study, trial; prune=true) |
| 116 | +else |
| 117 | + tell(study, trial, score) |
| 118 | +end |
| 119 | +``` |
| 120 | + |
| 121 | +Use `suggest_int` / `suggest_categorical` / `suggest_float` (with `log=true` for scale-free |
| 122 | +quantities like learning rates) to declare each dimension. The example searches architecture |
| 123 | +(`mps`, `layer_size`, `hidden_layers`), optimiser (`optimizer`, `lr`, `lr_decay_ratio`, |
| 124 | +`weight_decay`), regularisation (`noise_std`), normalisation (`norm_type`), and training strategy |
| 125 | +(`random_sampling`, `window_size`). |
| 126 | + |
| 127 | +### 4. Resume-awareness and results |
| 128 | + |
| 129 | +Because the study is persistent, the loop counts already-completed trials so restarts converge on a |
| 130 | +fixed total instead of adding a fresh batch each time: |
| 131 | + |
| 132 | +```julia |
| 133 | +n_completed = length(study.study.trials) |
| 134 | +n_remaining = max(0, n_trials - n_completed) |
| 135 | +``` |
| 136 | + |
| 137 | +When the run finishes, inspect the outcome with `best_trial(study)`, `best_params(study)`, and |
| 138 | +`best_value(study)`. |
| 139 | + |
| 140 | +## Adapting it to your own dataset |
| 141 | + |
| 142 | +1. **Point at your data** — set `ds_path`, and generate or provide `train.h5` / `valid.h5` / |
| 143 | + `test.h5` + `meta.json`. |
| 144 | +2. **Precompute normalization** — if you search `norm_type`, run [`update_meta!`](@ref GraphNetSim.update_meta!) |
| 145 | + once per statistic you want available (`:minmax` and/or `:meanstd`); a trial then only selects |
| 146 | + between them. |
| 147 | +3. **Set the fixed budget** — `n_steps` per trial, `cp_interval` (validation cadence), `n_trials`, |
| 148 | + and the simulation interval (`dt`, `tstop`). These trade search breadth against wall-clock time. |
| 149 | +4. **Edit the search space** — add or remove `suggest_*` calls, mirror them in the `params` |
| 150 | + NamedTuple, and forward them to `train_network`. |
| 151 | + |
| 152 | +!!! warning "Only tune parameters the API actually exposes" |
| 153 | + Every keyword forwarded to `train_network` must be an [`Args`](@ref GraphNetSim.Args) field, and |
| 154 | + every keyword to a strategy constructor must exist on that strategy (for example, |
| 155 | + [`DerivativeTraining`](@ref GraphNetSim.DerivativeTraining) accepts only `window_size` and |
| 156 | + `random`). Passing an unknown keyword errors when the trial builds its configuration. |
| 157 | + |
| 158 | +## Another instance |
| 159 | + |
| 160 | +[`example/WaterRamps/WaterRamps_optuna.jl`](https://github.com/una-auxme/GraphNetSim.jl/tree/main/example/WaterRamps/WaterRamps_optuna.jl) |
| 161 | +applies this same workflow to the (external) WaterRamps dataset, which is useful as a larger-scale |
| 162 | +reference — see [Examples](@ref) for why those research scripts are not turnkey. |
0 commit comments