MetNet-2 / README.md
zhangrenchao's picture
Publish MetNet-2 reproduction
efe4fbe verified
|
Raw
History Blame Contribute Delete
5.72 kB
---
license: apache-2.0
language:
- en
tags:
- OneScience
- earth-science
- precipitation-forecasting
- MetNet-2
frameworks: PyTorch
---
<p align="center">
<strong><span style="font-size: 30px;">MetNet-2</span></strong>
</p>
# Model Introduction
MetNet-2 generates high-resolution probabilistic precipitation forecasts up to 12 hours ahead from radar, satellite, and atmospheric-state inputs for short-range forecasting and extreme-precipitation risk analysis.
Paper: Deep learning for twelve hour precipitation forecasts
https://doi.org/10.1038/s41467-022-32483-x
# Model Description
The method was proposed by the Google Research team. The paper constructs its 2017-2020 training and test data from MRMS radar, GOES satellite imagery, and HRRR assimilated atmospheric states. The model predicts grid-cell precipitation probability distributions at two-minute intervals up to 12 hours ahead.
# Use Cases
| Use Case | Description |
| :---: | :--- |
| Categorical precipitation forecasting | Validate the 512-category conditional distribution and 12-hour lead-time protocol. |
| Core-method validation | Validate ConvLSTM, lead-time FiLM, and multiscale dilated residual stacks. |
| Local engineering validation | Exercise the complete logical `641×512×512` contract through a deterministic procedural field without materializing the full input. |
| ModelScope/OneCode execution | Validate structured data, training, inference, probabilistic precipitation metrics, and visualization in ModelScope or OneCode environments. |
| Multi-GPU training | Validate distributed training and the checkpoint workflow through `torchrun`. |
# Usage Instructions
## 1.OneCode
Experience intelligent, one-click AI4S programming through the OneCode online environment:
[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
## 2. Download and Installation
```bash
hf download OneScience-Group/MetNet-2 --local-dir ./MetNet-2
cd MetNet-2
```
### 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 in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
**DCU Environment**
```bash
# 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**
```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
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
```
### Training Data
The paper training inputs combine MRMS, HRRR, GOES, static geography, and temporal information under a 641-channel, `512×512` spatial-domain protocol. This repository creates eight deterministic window records and constructs only the selected `32×32` windows plus their halos at runtime while retaining the complete logical shape and channel grouping. The synthetic data validate engineering connectivity only and do not represent real meteorological distributions, paper-scale training, or paper performance.
```bash
python scripts/fake_data.py
```
### Training
For single-GPU training, use:
```bash
python scripts/train.py
```
For multi-GPU training, use:
```bash
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
```
The default reduces sample count, network width, residual blocks, and training steps without reducing the 641 channels, 512 precipitation categories, or 12-hour lead protocol. Formal experiments require real MRMS, GOES, and HRRR data and full computing resources, with artifacts saved to:
```text
result/checkpoints/metnet_2.pt
result/training/metrics.json
```
### Trained Weights
The paper does not provide directly loadable official model weights, and this repository bundles no weights under `weight/`. The locally trained checkpoint is saved to `result/checkpoints/metnet_2.pt` and must not be represented as an official pretrained weight.
### Inference
```bash
python scripts/inference.py
```
Inference loads the training checkpoint and generates 512-category precipitation probabilities and their CDF for selected spatial windows while preserving coverage and completeness metadata. Complete numerical results are saved to:
```text
result/output/predictions.npz
```
### Evaluation and Visualization
```bash
python scripts/result.py
```
Evaluation computes discrete CRPS, Brier Score, and CSI at multiple precipitation thresholds and generates target, expected-rate, and error comparisons. Synthetic-data results validate engineering only and do not represent paper performance; outputs are saved to:
```text
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 MetNet-2 specifications.
Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.