A sandbox repository containing instructions and notebooks to train, evaluate, and use ML models using Anemoi.
- Setup Environment
- Config Archive
- Explore Anemoi Documentation
- Example Training Run: GraphCast Global Model
- Run Lightweight Verification
Download and install miniforge, a minimal install for Conda and Mamba specific to conda-forge.
Create the environment:
conda create -n anemoi python=3.10 -y
conda activate anemoi
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121Developers:
pip install anemoi-utils
git clone git@github.com:ecmwf/anemoi-datasets.git
cd anemoi-datasets
pip install -e .
cd ..
git clone git@github.com:ecmwf/anemoi-core.git
cd anemoi-core
cd graphs && pip install -e . && cd ..
cd models && pip install -e . && cd ..
cd training && pip install -e . && cd ..
cd ..
git clone git@github.com:ecmwf/anemoi-inference.git
cd anemoi-inference
pip install -e .
cd ..The following environment has to be created in order to run training and inference with older versions of anemoi.
Create the environment:
conda create -n anemoi-legacy python=3.10 -y
conda activate anemoi-legacy
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install anemoi-datasets==0.5.23
pip install anemoi-graphs==0.4.4
pip install anemoi-inference==0.5.4
pip install anemoi-models==0.4.2
pip install anemoi-training==0.3.2
pip install anemoi-transform==0.1.10
pip install anemoi-utils==0.4.22
pip install eccodes-cosmo-resources-python
conda env config vars set ECCODES_DEFINITION_PATH=$CONDA_PREFIX/share/eccodes-cosmo-resources/definitionsAlternatively, refer to anemoi-config.
Anemoi relies on YAML configuration files to define pipelines.
You can find a full archive of MeteoSwiss use cases:
git clone git@github.com:MeteoSwiss/ml-pilot-anemoi.gitPlease refer to Hydra Introduction for more details about hydra configurations.
You can have a look at all available CLI command within anemoi in each documentation section:
Here is an example on how to setup a training run on a HPC using anemoi-training.
Download the ERA5 096 dataset from the anemoi-catalogue.
anemoi-datasets copy --resume s3://ml-datasets/aifs-ea-an-oper-0001-mars-o96-1979-2022-6h-v6.zarr .Update paths in config file: configs/stage_A-metno.yaml.
Create the bash script called stage_A-metno.sh to submit the job on the HPC machine.
stage_A-metno.sh:
#!/bin/bash
#SBATCH --job-name="stage_A"
#SBATCH --partition=normal
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --gpus-per-node=4
#SBATCH --cpus-per-task=16
#SBATCH --mem=256G
#SBATCH --time=24:00:00
#SBATCH --no-requeue
#SBATCH --exclusive
#SBATCH --output=stage_A-metno.log
#SBATCH --error=stage_A-metno.err
echo "Running training job on $(hostname) at $(pwd)"
# Activate environment
source /users/apennino/.bashrc
conda activate anemoi
srun -u anemoi-training train \
--config-dir="configs" \
--config-name="stage_A-metno"You can now launch the training job by:
sbatch stage_A-metno.shMonitor the training progress by checking the queue:
squeue -u YOUR-USERNAMEOr check the .log and .err files created by the script.
You can check the evolution of the training by checkig progress in the MLFlow server.
You can find your specific run by looking at the values of experiment_name and run_name.
Please refer to MLFlow Introduction for more details about MLFlow monitoring.
Once your training is completed, anemoi generates a checkpoint file containing all model weights. You can find it by looking at the checkpoint path in the config file.
Until it is moved to its own repository, you can find the verification tool here:
git clone git@github.com:MeteoSwiss/ml-pilot-anemoi.git
cd ml-pilot-anemoi
git checkout ml-forecasting
cd inference
-
Locate your
run_idin MLflow.
-
Launch:
./submit_inference_array.sh 090030db9b0c4d3c8e9f42ef66053e6b -m stage_A-metnoverif/: metrics vs. analysis dataplots/: images and animations
grib/: GRIB files for NWPraw/: raw NumPy predictions
Execute the verification.ipynb notebook to check the desired metrics and models.