Modelpublic

OneScience-Group/GenCast

sourceHugging Faceapache-2.0updated 1mo agoView on Hugging Face

0likes5downloads

Model Card

GenCast

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

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 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 --repo-type 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.