OneScience-Group/FourCastNet_v2

<p align="center"> <strong> <span style="font-size: 30px;">FourCastNet_v2</span> </strong> </p>

Model Introduction

FourCastNet v2 is a global weather forecast model based on the Spherical Fourier Neural Operator (SFNO), proposed by NVIDIA and its collaborators.

Paper: Spherical Fourier Neural Operators: Learning Stable Dynamics on the Sphere

https://arxiv.org/abs/2306.03838

Model Description

The key architectural change from v1 is replacing the Adaptive Fourier Neural Operator (AFNO) with the Spherical Fourier Neural Operator (SFNO).

Use Cases

Usage Guide

1. OneCode Usage

Experience intelligent one-click AI4S programming through the OneCode online environment:

Click to Experience Intelligent One-Click AI4S Programming

2. Manual Installation and Usage

Hardware Requirements

• —A GPU or DCU is recommended.
• —CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
• —DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.

Download the Model Package

hf download OneScience-Group/FourCastNet_v2 --local-dir ./FourCastNet_v2
cd FourCastNet_v2

Install the Runtime Environment

DCU Environment

# Please activate DTK and CONDA first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# uv installation is supported
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

GPU Environment

# Please 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
# uv installation is supported
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

Training Data Introduction

The OneScience community provides an ERA5 data slice that can be downloaded as follows:

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

Real HDF5 annual files must contain fields, variable attributes, time_step, global_means, and global_stds. When real data is unavailable, first generate synthetic files for pipeline validation:

python scripts/fake_data.py

Training

Single GPU:

python scripts/train.py

Multi-GPU:

torchrun --nproc_per_node=8 scripts/train.py

The default checkpoint is saved to data/checkpoint/one_step/model_bak.pt.

Fine-tuning

Single GPU:

python scripts/train.py --stage finetune

Multi-GPU:

torchrun --nproc_per_node=8 scripts/train.py --stage finetune

The checkpoint is saved to data/checkpoint/<stage>/model_bak.pt by default.

Training Weights

This repository provides weights trained on ERA5 reanalysis data in the weight/ folder. The weight files will be uploaded soon and are expected to be available in the near future.

Inference

python scripts/inference.py

Prediction results are written to result/output/ by default.

Evaluation and Visualization

python scripts/result.py

The default output includes latitude-weighted RMSE/ACC metrics and result/figures/t2m_forecast.png.

Official OneScience Resources

Citation and License

• —The SFNO numerical implementation of FourCastNet v2 follows the design of NVIDIA Earth2MIP and related official implementations. The upstream code and model licenses and copyright notices must be retained.