| --- |
| datasets: |
| - OneScience/ERA5 |
| frameworks: |
| - "" |
| language: |
| - en |
| license: mit |
| tags: |
| - OneScience |
| - Earth Science |
| - ERA5 |
| - Medium-Range Weather Forecasting |
| - ViT |
| tasks: [] |
| --- |
| <p align="center"> |
| <strong> |
| <span style="font-size: 30px;">Stormer</span> |
| </strong> |
| </p> |
| |
|
|
| # 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](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/Stormer --local-dir ./Stormer |
| cd Stormer |
| ``` |
|
|
| ### 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 |
| python scripts/train.py |
| ``` |
|
|
| Multiple GPUs: |
|
|
| ```bash |
| 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 |
|
|
| ```bash |
| python scripts/inference.py |
| ``` |
|
|
| By default, inference loads `data/checkpoints/model_bak.pth`, and results are saved to `result/output/`. |
|
|
|
|
| ### 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 Stormer paper. |
|
|
