| --- |
| frameworks: JAX |
| language: |
| - en |
| license: apache-2.0 |
| tags: |
| - OneScience |
| - Earth Science |
| - Weather Forecasting |
| - Ensemble Forecasting |
| - ERA5 |
| tasks: [] |
| datasets: |
| - OneScience/ERA5 |
| --- |
| |
| <p align="center"> |
| <strong> |
| <span style="font-size: 30px;">GenCast</span> |
| </strong> |
| </p> |
| |
| # Model Overview |
|
|
| GenCast is a probabilistic global weather forecasting model developed by Google DeepMind. Its paper appeared as the cover article of the leading scientific journal *Nature* on December 4, 2024. |
|
|
| Paper: *GenCast: Diffusion-Based Ensemble Forecasting for Medium-Range Weather* |
|
|
| https://arxiv.org/abs/2312.15796 |
|
|
| # Model Description |
|
|
| GenCast is an ensemble forecasting model built with graph neural networks and diffusion models. Across a comprehensive set of evaluations, it outperformed ENS, the European Centre for Medium-Range Weather Forecasts' (ECMWF) leading ensemble forecasting system. |
|
|
|
|
| # Use Cases |
|
|
| | Use Case | Description | |
| | :---: | :--- | |
| | Weather forecasting training | Train the model on ERA5 data in HDF5 format that conforms to the GenCast data protocol. | |
| | Quick local validation | Use synthetic data to validate data loading, model training and inference, and visualization of inference results. | |
| | ModelScope/OneCode execution | Download the standalone model package, install its dependencies, and run the included scripts directly. | |
| | Multi-GPU training | Use JAX `pmap` for data-parallel training across multiple GPUs or accelerators on a single host. | |
|
|
|
|
| # Usage |
|
|
| ## 1. Using OneCode |
|
|
| Use the OneCode online environment for an intelligent, one-click AI4S development experience: |
|
|
| [Try one-click AI4S development with OneCode](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home) |
|
|
| ## 2. Manual Setup |
|
|
| **Hardware Requirements** |
|
|
| - A GPU or DCU is recommended. |
| - A CPU can be used for import checks and connectivity validation with a minimal configuration, but full training and inference will be slow. |
| - DCU users must install DTK in advance. DTK 25.04.2 or later is recommended; alternatively, use the OneScience-recommended version compatible with your cluster. |
|
|
| ### Download the Model Package |
|
|
| ```bash |
| hf download --model OneScience-Group/GenCast --local-dir ./GenCast |
| cd GenCast |
| ``` |
|
|
| ### Set Up the Runtime Environment |
|
|
| **DCU Environment** |
|
|
| ```bash |
| # Activate DTK and conda first. |
| conda create -n onescience311 python=3.11 -y |
| conda activate onescience311 |
| # Installation with uv is also supported. |
| pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai |
| ``` |
|
|
| **GPU Environment** |
|
|
| ```bash |
| # Activate conda first. |
| conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12 |
| conda activate onescience311 |
| # Installation with uv is also supported. |
| pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai |
| ``` |
|
|
| ### Training Data |
|
|
| The OneScience community provides ERA5 data for training. Because of file-size constraints, the repository currently contains a self-contained data slice. Download the data with the following command and ensure that the data path in `conf/config.yaml` is configured correctly: |
|
|
| ```bash |
| hf download --dataset OneScience-Group/ERA5 --local-dir ./data |
| ``` |
|
|
|
|
| ### Training |
|
|
| Single GPU: |
|
|
| ```bash |
| # If real data is unavailable, first run `python scripts/fake_data.py` to generate synthetic data. |
| python scripts/train.py |
| ``` |
|
|
| Multiple GPUs: |
|
|
| ```bash |
| CUDA_VISIBLE_DEVICES=0,1 python scripts/train.py --config conf/config.yaml --parallel-mode pmap --num-devices 2 --global-batch-size 2 |
| # CUDA_VISIBLE_DEVICES specifies the GPU indices to expose. |
| # --num-devices specifies the number of GPUs to use. |
| # --global-batch-size specifies the batch size and must be divisible by the number of GPUs. |
| ``` |
|
|
| After training, the weights are saved to `data/checkpoints/model_bak.npz`. |
|
|
|
|
| ### Pre-trained Weights |
|
|
| This repository will provide weights trained on ERA5 reanalysis data in the `weights/` directory. The weight files are being prepared and will be uploaded soon. |
|
|
|
|
| ### Inference |
|
|
| By default, inference loads `data/checkpoints/model_bak.npz`: |
|
|
| ```bash |
| python scripts/inference.py |
| ``` |
|
|
| ### Evaluation and Visualization |
|
|
| ```bash |
| python scripts/result.py |
| ``` |
|
|
|
|
| # Official OneScience Resources |
|
|
| | Platform | OneScience Main Repository | Skills Repository | |
| | --- | --- | --- | |
| | Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills | |
| | GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills | |
|
|
|
|
| # Citation and License |
|
|
| - This repository is a reproduction of the original GenCast paper. |
|
|
