FengWu-W2S

Model Introduction

FengWu-W2S (FengWu Weather-to-Subseasonal) is a seamless global weather-to-subseasonal forecasting model extending FengWu.

Paper: FengWu-W2S: A deep learning model for seamless weather-to-subseasonal forecast of global atmosphere

https://arxiv.org/abs/2411.10191

Model Description

The model uses a six-hour time step for autoregressive forecasts of up to 42 days. Coupled atmospheric, ocean, and land branches, together with diversity perturbations, are used to improve extended-range forecast skill. This repository contains a compact implementation of the coupled interfaces for reproducible workflow checks.

Use Cases

Scenario Description
Weather-to-subseasonal forecast research Train a 78-channel coupled model with six-hourly global atmospheric, ocean, and land fields.
Local quick validation Use small-grid synthetic HDF5 data to check training, fine-tuning, inference, and forecast visualization.
ModelScope / OneCode execution Download the standalone model package, install dependencies, and run the scripts directly.
Multi-GPU training Launch PyTorch DistributedDataParallel 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/FengWu-W2S --local-dir ./FengWu-W2S
cd FengWu-W2S

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 for training. Download it and confirm that the data path in conf/config.yaml is correct:

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

Real files must contain all 78 channels listed in conf/config.yaml, including atmospheric, ocean, and land variables, with six-hour time spacing. The configured group indices define the coupled branches.

Generate Synthetic Data

The default configuration describes a full 721x1440 grid. For a practical local smoke test, generate a small fixture explicitly:

python scripts/fake_data.py --height 32 --width 64 --timesteps 12

Synthetic HDF5 files are for workflow validation only and do not represent the ERA5 reanalysis or the paper's forecast quality.

Training

Single GPU:

python scripts/train.py

Multi-GPU:

torchrun --nproc_per_node=2 scripts/train.py

For a short smoke run, reduce the workload while keeping the generated data and model grid consistent:

python scripts/train.py --max-epoch 1 --batch-size 1 --num-workers 0 --rollout-steps 1

The default checkpoint is saved to data/checkpoints/model_bak.pth.

Fine-tuning

Resume from the base checkpoint with the lower fine-tuning learning rate:

python scripts/train.py --finetune

Training Weights

This repository provides a weight/ directory for FengWu-W2S checkpoints. The weight files will be uploaded soon and are expected to be available in the near future.

Inference

Inference reads data/checkpoints/model_bak.pth by default and writes forecasts grouped by initialization year to result/output/<year>/, together with result/output/index.json:

python scripts/inference.py

Use --steps and --limit to bound a local smoke test; --stochastic enables perturbation sampling.

Evaluation and Visualization

python scripts/result.py

The script computes per-channel RMSE, normalized RMSE, and anomaly ACC, and generates forecast-comparison, lead-time skill, channel-ranking, and training-loss figures.

Official OneScience Resources

Citation and License

  • Paper: https://arxiv.org/abs/2411.10191
  • This directory is an independent reproduction of the paper method and does not represent official code, weights, or training results released by the authors.
Downloads last month
12
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Paper for OneScience-Group/FengWu-W2S