Lightning ERA5 U-Net Template

这是一个面向气象场建模的工程模板，目标是让你学习一个较完整的深度学习项目范式：

• PyTorch Lightning 负责训练框架
• wandb 负责实验记录
• setup.py 负责项目打包
• src 目录布局负责工程边界
• U-Net 负责二维空间场建模
• ERA5 .nc 日频文件负责数据输入

任务定义

这个模板默认将问题定义为：

• 输入：过去 history_days 天的 ERA5 多变量二维场
• 输出：未来 lead_days 天后的目标变量二维场

一个典型例子是：

• 输入变量：["t2m", "u10", "v10", "msl"]
• 目标变量：["t2m"]
• history_days=5
• lead_days=1

即用过去 5 天的多个变量，预测下一天的 2 米气温场。

目录结构

application/
├─ configs/
│  └─ train_unet.yaml
├─ scripts/
│  └─ train.py
├─ src/
│  └─ lightning_unet/
│     ├─ data/
│     │  ├─ datamodule.py
│     │  └─ era5.py
│     ├─ models/
│     │  └─ unet.py
│     ├─ modules/
│     │  └─ forecast.py
│     ├─ utils/
│     │  └─ config.py
│     └─ __init__.py
├─ requirements.txt
├─ setup.py
└─ .gitignore

数据组织约定

当前模板假设：

1. data_dir 下每个 .nc 文件对应一天。
2. 文件名按字典序排序后，就是时间顺序。
3. 每个文件里都包含同样的变量，比如 t2m/u10/v10/msl。
4. 每个文件里的空间网格一致。
5. 变量至少带有 lat/lon 或 latitude/longitude 这样的二维空间维。
6. 如果文件中仍有 time 维，模板会默认取第一个时间片。

例如：

data/era5_daily/
├─ era5_20200101.nc
├─ era5_20200102.nc
├─ era5_20200103.nc
└─ ...

配置说明

关键配置位于 configs/train_unet.yaml：

• data.data_dir: ERA5 日频 .nc 文件目录
• data.input_variables: 输入变量名列表
• data.target_variables: 预测目标变量名列表
• data.history_days: 输入历史窗口长度
• data.lead_days: 预测提前期
• data.train_split: 训练集比例
• data.val_split: 验证集比例
• data.normalization: 每个变量的均值和标准差

快速开始

cd application
python -m venv .venv
.venv\Scripts\activate
pip install -r requirements.txt
pip install -e .
python scripts/train.py --config configs/train_unet.yaml

wandb

本模板默认启用 wandb。如果只想本地记录：

set WANDB_MODE=offline
python scripts/train.py --config configs/train_unet.yaml

当前实现的工程职责

• models/unet.py 只负责网络结构
• modules/forecast.py 负责训练、验证、loss 和 metric
• data/era5.py 负责 ERA5 .nc 文件读取与样本切片
• data/datamodule.py 负责 DataLoader 和数据集切分
• scripts/train.py 负责组装配置、logger、trainer 和 callback

重要假设

这个模板是教学型模板，因此做了几个刻意简化：

1. 默认使用 U-Net 预测二维格点场。
2. 默认把时间历史维展开到 channel 维。
3. 默认按文件顺序构造滑动窗口样本。
4. 默认使用配置中提供的均值和标准差做归一化。
5. 默认用 MSE + MAE 作为主要训练和评估指标。

你后续可以继续扩展

1. 增加 test 集与独立推理脚本。
2. 增加更多 meteorology metric，比如 anomaly correlation、RMSE、bias。
3. 把单步预测扩展成多步预测。
4. 加入 mask 机制处理陆海或缺测区域。
5. 把 YAML 配置升级成 Hydra 或 OmegaConf。