OneScience-Group/ClimaX

ClimaX

Model Overview

ClimaX is the first general-purpose foundation model designed for weather and climate science. It was jointly developed by researchers at Microsoft Research and the University of California, Los Angeles (UCLA), and was presented at ICML 2023.

Paper: ClimaX: A Foundation Model for Weather and Climate

https://arxiv.org/abs/2301.10343

Model Description

ClimaX is built on the Vision Transformer architecture. It was the first model to unify pre-training on heterogeneous climate datasets and adaptation to downstream tasks, such as regional forecasting and climate projections, within a single framework.

Use Cases

Use Case	Description
Weather forecasting training	Train ClimaX on ERA5 data in HDF5 format.
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	Launch multi-process training with torchrun.

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

hf download --model OneScience-Group/ClimaX --local-dir ./ClimaX
cd ClimaX

Set Up the Runtime Environment

DCU Environment

# 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

# 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:

hf download --dataset OneScience-Group/ERA5 --local-dir ./data

Training

Single GPU:

python scripts/train.py

Multiple GPUs:

torchrun --nproc_per_node=8 scripts/train.py

Training saves the model_bak.pth checkpoint under data/checkpoints/.

Pre-trained Weights

This repository will provide weights trained on ERA5 data in the weights/ directory. The weight files are being prepared and will be uploaded soon.

Inference

python scripts/inference.py

Inference results are saved to result/output/.

Evaluation and Visualization

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 ClimaX paper and is licensed under the MIT License.