license: cc-by-nc-nd-4.0
language:
- en
tags:
- OneScience
- Earth Science
- Rainfall Correction
- UnetDif
frameworks: PyTorch
UnetDif
Model Introduction
UnetDif uses the paper's bias-targeted multitask U-Net to correct short-range heavy rainfall forecasts.
Paper: Bias-targeted deep learning enhances short-range heavy rainfall forecasts
https://doi.org/10.1038/s41612-026-01366-z
Model Description
The method reproduced by UnetDif was proposed by research teams from Zhejiang Meteorological Observatory, the National Meteorological Center, and Zhejiang University. The paper constructs Yangtze River Delta samples from European Centre for Medium-Range Weather Forecasts (ECMWF) fields, China Meteorological Administration Multi-source Precipitation Analysis (CMPA), and terrain data. The model learns ECMWF - CMPA rainfall bias and corrects three-hour precipitation through a shared U-Net backbone with four heads for dry areas, false alarms, positive bias, and negative bias.
Use Cases
| Use Case | Description |
|---|---|
| Three-hour rainfall-bias correction | Generate corrected ECMWF rainfall from 39-channel inputs on the fixed Yangtze River Delta grid. |
| ModelScope/OneCode execution | Validate synthetic-data generation, six-loss training, inference, evaluation, and visualization in ModelScope or OneCode environments. |
| Hugging Face/OneCode execution | Run the complete engineering workflow in Hugging Face or OneCode environments. |
| Multi-GPU training | Launch distributed data-parallel training with torchrun. |
Usage Instructions
1.OneCode
Experience intelligent, one-click AI4S programming through the OneCode online environment:
Try intelligent, one-click AI4S programming
2. Download and Installation
hf download OneScience-Group/UnetDif --local-dir ./UnetDif
cd UnetDif
Environment Dependencies
Hardware Requirements
- A GPU or DCU is recommended.
- A CPU can be used for connectivity validation with the default small-sample configuration.
- DCU users must install DTK first. DTK 25.04.2 or later, or the OneScience-recommended version matching the cluster, is recommended.
DCU Environment
# Activate DTK and Conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
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
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
Training Data
Training data contain ECMWF weather forecasts, CMPA precipitation observations, and static terrain. Each sample uses 39 input channels, with corresponding CMPA precipitation and ECMWF precipitation bias as targets. The data use a 56×56 Yangtze River Delta grid at 0.125° resolution. Each rain day contains eight consecutive three-hour periods, with precipitation in mm/3h. Synthetic data validate training, inference, and evaluation only and do not represent the distributions of ECMWF, CMPA, real terrain, or formal paper performance.
python scripts/fake_data.py
Training
For single-device training, use:
python scripts/train.py
For multi-GPU training, use:
torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
Training learns dry-area, false-alarm, and signed precipitation-bias tasks and saves the model checkpoint and training metrics. Results are saved to:
result/checkpoints/unetdif.pt
result/training/metrics.json
Trained Weights
This repository does not include weights under weight/. The paper does not provide a confirmed official checkpoint, and the engineering checkpoint trained on synthetic data is for workflow validation only and is not claimed to be compatible with external weights.
Inference
python scripts/inference.py
Inference loads the checkpoint and applies all four output heads under bias = ECMWF - CMPA to produce nonnegative corrected rainfall. Results are saved to:
result/output/predictions.npz
Evaluation and Visualization
python scripts/result.py
Evaluation produces categorical and spatial-skill metrics for individual three-hour periods and 24-hour accumulated precipitation and saves structured results. It also generates comparisons of ECMWF, CMPA, corrected precipitation, and errors. Synthetic-data results validate the engineering workflow only and do not represent formal paper performance.
result/evaluation/metrics.json
result/evaluation/comparison.png
Official OneScience Information
| 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 an independent engineering reproduction of the public UnetDif specifications.
Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.