DiffCast-CVPR2024

Official implementation of "DiffCast: A Unified Framework via Residual Diffusion for Precipitation Nowcasting"

Introduction

DiffCast is a precipitation nowcasting framework based on diffusion model and a deterministic predictive backbone, which can be achieved with various spatio-temporal predictive models optimized with deterministic loss (e.g., SimVP, Earthformer, ConvGRU, PhyDNet et al).

This repository contains the part of training and inference code for using DiffCast to make predictions (5 --> 20) on SEVIR datasets.

Code

Environment

conda env create -f env.ymal
conda activate diffcast

Optional Accelerate Env

We apply the HuggingFace Accelerator in our code to utilize multi-gpus. One can config the accelerator env before runing code.

• config the accelerate: accelerate config
• apply accelerate to run code: accelerate launch *.py

Resource

pretrained DiffCast_PhyDNet: GoogleDrive

Datasets

All the four datasets in our paper is publicly available. You can find the datasets as follows:

• SEVIR
• Metnet
• Shanghai_Radar
• CIKM_Radar

We update the preprocess code at HERE

Also, you can directly download the h5 file we have built:

• Metnet: GoogleDrive
• Shanghai_Radar: GoogleDrive

We apologize for the late update of our datasets for some reason, feel free to concat me ([email protected],[email protected]) if you have any questions about code or datasets.

Toy prediction visiualization

We give some demos from SEVIR to easily check the predictive performance of DiffCast_PhyDNet. Before that, you need to download the pretrained checkpoint and put it in resources/

python sample_batch.py

Evaluation

# Note: Config the dataset path in `dataset/get_dataset.py` before running.
python run.py --backbone phydnet --use_diff --eval --ckpt_milestone resources/diffcast_phydnet_sevir128.pt  

Backbone Training

python run.py --backbone simvp

You can check the experimental configuration by

python run.py -h

Acknowledgement

We refer to implementations of the following repositories and sincerely thank their contribution for the community:

• denoising_diffusion_pytorch
• Prediff
• RVD
• OpenSTL

Citation

@inproceedings{Yu2024diffcast,
  title={DiffCast: A Unified Framework via Residual Diffusion for Precipitation Nowcasting},
  author={Demin Yu and Xutao Li and Yunming Ye and Baoquan Zhang and Chuyao Luo and Kuai Dai and Rui Wang and Xunlai Chen},
  booktitle={The IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  year={2024}
}