This example is a lightweight implementation of the DeepMind’s GraphCast model in PyTorch, to provide a recipe on how to train such GNN models in Modulus.
GraphCast is a multi-scale graph neural network-based autoregressive model. It is trained on historical weather data from ECMWF’s ERA5 reanalysis archive. GraphCast generates predictions at 6-hour time steps for a set of surface and atmospheric variables. This prediction covers a 0.25-degree latitude-longitude grid, providing approximately 25 x 25 kilometer resolution at the equator.
The model is trained on a 73-channel subset of the ERA5 reanalysis data on single levels and pressure levels that are pre-processed and stored into HDF5 files. A 20-channel subset of the ERA5 training data is hosted at the National Energy Research Scientific Computing Center (NERSC). For convenience it is available to all via Globus. You will need a Globus account and will need to be logged in to your account in order to access the data. You may also need the Globus Connect to transfer data.
Please refer to the reference paper to learn about the model architecture.
To train the model on a single GPU, run
python train_graphcast.py
This will launch a GraphCast training with up to 12 steps of fine-tuning
using the base configs specified in the constants.py file.
Data parallelism is also supported with multi-GPU runs. To launch a multi-GPU training, run
mpirun -np <num_GPUs> python train_graphcast.py
If running in a docker container, you may need to include the
--allow-run-as-root in the multi-GPU run command.
Progress and loss logs can be monitored using Weights & Biases. This requires to have an active Weights & Biases account. You also need to provide your API key. There are multiple ways for providing the API key but you can simply export it as an environment variable
export WANDB_API_KEY=<your_api_key>
The URL to the dashboard will be displayed in the terminal after the run is launched.
If needed, Weights & Biases can be disabled by
export WANDB_MODE='disabled'