Merge ACE2.1-ERA5 (AIMIP) training and evaluation baseline configs#1027
Merge ACE2.1-ERA5 (AIMIP) training and evaluation baseline configs#1027
Conversation
Remove intermediate fine-tuning explorations superseded by the separate-decoder + LR warmup approach (the final model). Trim the fine-tuning and seed-selection launch scripts to reference only the final approach. Deleted configs (non-final fine-tuning variants): - ace-fine-tune-decoder-pressure-level-config.yaml - ace-fine-tune-decoder-pressure-level-lr-warmup-config.yaml - ace-fine-tune-decoder-pressure-level-frozen-config.yaml - ace-fine-tune-decoder-pressure-level-frozen-lr-warmup-config.yaml - ace-fine-tune-decoder-pressure-level-reweight-config.yaml - ace-fine-tune-decoder-pressure-level-separate-decoder-config.yaml - restart-ace-fine-tune-decoder-pressure-levels.sh
Resolves conflicts in fme/core/step/single_module.py and test_step.py by accepting main's SecondaryDecoderConfig/SecondaryDecoder approach and dropping the branch's inline MLP + additional_diagnostic_names approach. Updates AIMIP configs to use the new secondary_decoder config format and moves loss/parameter_init from stepper to stepper_training per the TrainConfig restructuring in main.
Delete 15 pre-generated IC-specific config files and instead do the _r[N]i label substitution inside the gantry container at job runtime via sed, keeping only the 3 template configs committed.
configs/baselines/era5-aimip/ace-evaluator-seed-selection-single-config.yaml
Show resolved
Hide resolved
Arcomano1234
left a comment
Arcomano1234
left a comment
There was a problem hiding this comment.
Left a few comments / questions for my own curiosity but this looks mostly good to go. I've been using some of these scripts so it will be nice to have in main. My only real comment is removing your hard-coded wandb name in a lot of the job submission scripts.
| - `run-ace-evaluator-seed-selection-single.sh` — single continuous 36-year run (1978-10-01 to | ||
| 2014-12-31). Config: `ace-evaluator-seed-selection-single-config.yaml`. | ||
|
|
||
| After reviewing results, update the base checkpoint ID in `run-ace-fine-tune-decoder-pressure-levels.sh`. |
There was a problem hiding this comment.
Question: What criteria did we use for AIMIP to determine the best seed? As I am still not sure I understand what the run-ace-evaluator-seed-selection.sh eval tells use compared to run-ace-evaluator-seed-selection-single.sh.
My suggestion is maybe add one more sentence on how the best seed was chosen at this stage.
There was a problem hiding this comment.
Will add a sentence, something like "hopefully the seed with the best time-mean climate/trends is the same across both evaluations otherwise make a subjective decision".
There was a problem hiding this comment.
Thanks, its definitely the most hand-wavy part of the process and even something vaguely similar to what we did is fine just so its documented
| CONFIG_FILENAME="ace-evaluator-seed-selection-single-config.yaml" | ||
| SCRIPT_PATH=$(git rev-parse --show-prefix) # relative to the root of the repository | ||
| CONFIG_PATH=$SCRIPT_PATH/$CONFIG_FILENAME | ||
| BEAKER_USERNAME=bhenn1983 |
There was a problem hiding this comment.
Nit: I didn't check all of the bash scripts but it would be nice to not have your username hardcoded here. I suggest BEAKER_USERNAME=$(beaker account whoami --format=json | jq -r '.[0].name')
There was a problem hiding this comment.
Thanks for catching this. Since my beaker name and wandb name don't match I always have to change this 🙄
2 participants
This PR adds the full set of scripts and configurations for "ACE2.1-ERA5" — a modification of the deterministic ACE2-ERA5 model trained and evaluated under the AIMIP protocol. It also merges
mainto pick up theSecondaryDecoderConfigAPI used for the pressure-level decoder fine-tuning stage, while retaining all of the job names, output paths, and checkpoint IDs as used on the original branch where the workflow actually occurred.Changes:
configs/baselines/era5-aimip/— new directory containing all scripts and configs for the ACE2.1-ERA5 pipeline (previouslyconfigs/baselines/era5/aimip/)run-ace-train.sh/ace-train-config.yaml— train 4-seed ensemble on ERA5 1979–2008run-ace-evaluator-seed-selection.sh/run-ace-evaluator-seed-selection-single.sh— evaluate trained and fine-tuned checkpoints to select best seedsrun-ace-fine-tune-decoder-pressure-levels.sh/ace-fine-tune-pressure-level-separate-decoder-config.yaml— fine-tune a secondary MLP decoder for 65 pressure-level diagnostic variables, usingsecondary_decoder(main'sSecondaryDecoderConfig)run-ace-inference.sh/ace-aimip-inference-{,p2k-,p4k-}config.yaml— 46-year inference with 5 ICs × 3 SST scenarios; IC label expansion done via inlinesedat job time (eliminates 15 near-identical committed config files)README.md— documents the intended workflowTests added
If dependencies changed, "deps only" image rebuilt and "latest_deps_only_image.txt" file updated