OneScience-Group/Stormer

Stormer

Model Overview

Stormer was jointly developed by researchers at Argonne National Laboratory and the University of California, Los Angeles (UCLA). Its core paper was published at NeurIPS 2024, a leading conference in artificial intelligence.

Paper: Scaling Transformer Neural Networks for Skillful and Reliable Medium-Range Weather Forecasting

https://arxiv.org/abs/2312.03876

Model Description

Stormer uses a standard Vision Transformer architecture and provides a streamlined deep learning model for medium-range weather forecasting.

Use Cases

Use Case	Description
Weather forecasting training	Train Stormer 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/Stormer --local-dir ./Stormer
cd Stormer

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 reanalysis data in the weights/ directory. The weight files are being prepared and will be uploaded soon.

Inference

python scripts/inference.py

By default, inference loads data/checkpoints/model_bak.pth, and 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 Stormer paper.