简介

本仓库实现了一个基于 PyTorch 的编码器-预报器模型，使用了 RNN（包括 TrajGRU 和 ConvLSTM）来进行降水临近预报。有关 TrajGRU 的更多信息，请参阅 HKO-7。

如果您对我的 ConvLSTM 和 TrajGRU 实现感兴趣，请查看 ConvLSTM 和 TrajGRU。假设输入的形状应为 。我的所有实现都在 HKO-7 数据集中被证明是有效的。希望这对您的研究有所帮助。

训练

首先，您需要从 HKO-7 申请 HKO-7 数据集，并修改 config.py 中的相关行以指定数据集路径。 其次，运行 python3 experiments/trajGRU_balanced_mse_mae/main.py，然后运行 python3 experiments/trajGRU_frame_weighted_mse/main.py，因为我是在上一步训练好的模型基础上进行了微调。

环境

Python 3.6+、PyTorch 1.0 以及 Ubuntu 或 macOS。

演示

性能

在 HKO-7 数据集上的性能如下所示。

CSI					HSS					平衡均方误差	平衡平均绝对误差
0.5496	0.4772	0.3774	0.2863	0.1794	0.6713	0.6150	0.5226	0.4253	0.2919	5860.97	15062.46

下载

Dropbox 预训练模型

引用

@inproceedings{xingjian2017deep,
    title={深度学习用于降水临近预报：基准与新模型},
    author={Shi, Xingjian 和 Gao, Zhihan 和 Lausen, Leonard 和 Wang, Hao 和 Yeung, Dit-Yan 和 Wong, Wai-kin 和 Woo, Wang-chun},
    booktitle={神经信息处理系统进展},
    year={2017}
}
@inproceedings{xingjian2015convolutional,
  title={卷积 LSTM 网络：一种用于降水临近预报的机器学习方法},
  author={Xingjian, SHI 和 Chen, Zhourong 和 Wang, Hao 和 Yeung, Dit-Yan 和 Wong, Wai-Kin 和 Woo, Wang-chun},
  booktitle={神经信息处理系统进展},
  pages={802--810},
  year={2015}
}

Precipitation-Nowcasting 快速上手指南

本指南帮助开发者快速部署基于 PyTorch 的降水临近预报模型（支持 TrajGRU 和 ConvLSTM），并在 HKO-7 数据集上进行训练。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Ubuntu 或 macOS
• Python 版本：3.6 及以上
• 深度学习框架：PyTorch 1.0 及以上
• 数据集：需提前申请 HKO-7 数据集

提示：国内用户安装 PyTorch 时，推荐使用清华或中科大镜像源以加速下载： pip install torch -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆项目代码

   git clone https://github.com/Hzzone/Precipitation-Nowcasting.git
   cd Precipitation-Nowcasting
   

2. 安装依赖 确保已安装 PyTorch，其他依赖通常包含在标准 Python 库中。如有 requirements.txt，可运行：

   pip install -r requirements.txt
   

3. 配置数据集路径 下载并解压 HKO-7 数据集后，打开项目根目录下的 config.py 文件，修改相关行以指向您的数据集实际存储路径。

基本使用

本项目采用两阶段训练策略：先进行基础训练，再在此基础上进行微调。

1. 第一阶段：基础模型训练 运行以下命令开始训练基础模型（使用 Balanced MSE/MAE 损失）：

   python3 experiments/trajGRU_balanced_mse_mae/main.py
   

2. 第二阶段：模型微调 在第一阶段生成的模型基础上，运行以下命令进行帧加权微调（Frame Weighted MSE）：

   python3 experiments/trajGRU_frame_weighted_mse/main.py
   

3. 使用预训练模型（可选） 如果仅需推理或测试，可直接下载预训练模型：

   • 下载地址：Dropbox Pretrained Model
   • 输入数据形状要求：$S \times B \times C \times H \times W$ （序列长度 × 批次大小 × 通道数 × 高度 × 宽度）