FuXi_v21

Model Introduction

FuXi 2.1 is a global deterministic machine-learning weather forecast model developed by Fudan University in collaboration with the Shanghai Artificial Intelligence Laboratory (SAIS). Its theoretical basis remains the original FuXi paper.

Paper: FuXi: A cascade machine learning forecasting system for 15-day global weather forecast

https://arxiv.org/abs/2306.12873

Model Description

The model addresses the excessive smoothing often observed in AI weather forecasts. It aims to produce clearer and more detailed forecast fields, improving the detection of extreme events such as heavy precipitation and strong winds without degrading conventional metrics such as root mean square error (RMSE).

Use Cases

Scenario Description
Global weather forecast training Train FuXi v2.1 with C85 ERA5 data in HDF5 format.
Local quick validation Use synthetic HDF5 data to check data loading, training, inference, and visualization of inference results.
ModelScope / OneCode execution Download the standalone model package, install dependencies, and run the scripts directly.
Multi-GPU training Run distributed data-parallel training with torchrun.

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/FuXi_v21 --local-dir ./FuXi_v21
cd FuXi_v21

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 training entry point uses the OneScience ERA5Dataset. The data root is specified by paths.data_root in conf/config.yaml. The OneScience community provides a data slice for interface validation and training:

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

Generate Synthetic Data

Real HDF5 files must contain a fields dataset, C85 variable attributes, six-hour intervals, and normalization statistics. When real data is unavailable, generate protocol-compatible synthetic files:

python scripts/fake_data.py

Training

Single GPU:

python scripts/train.py

Multi-GPU:

torchrun --nproc_per_node=8 scripts/train.py

Training checkpoints are saved to data/checkpoint/model_bak.pth by default, and metrics are saved to output/training/metrics.json.

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

Inference results are saved to output/inference/forecast.nc by default.

Evaluation and Visualization

python scripts/result.py

The default output is figures/fuxi21_t2m.png.

Official OneScience Resources

Citation and License

  • This project is an unofficial forward-graph reproduction of FuXi v2.1. It does not represent official weights or training recipes released by Fudan University.
  • This adapted repository is distributed under Apache License 2.0 metadata. ERA5 data, OneScience, and the upstream FuXi implementation remain subject to their respective official licenses and terms of use.
Downloads last month
10
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Paper for OneScience-Group/FuXi_v21