三亩地 三亩地SAN MU DI · CODE DIARY
ARTICLE DETAIL

日记详情

真实记录编程学习的某一天,欢迎挑你感兴趣的翻一翻。

ClimaX进阶实践:完整复现论文中的全球天气预报实验

ClimaX进阶实践:完整复现论文中的全球天气预报实验

ClimaX进阶实践:完整复现论文中的全球天气预报实验

【免费下载链接】ClimaXFoundation model for weather & climate项目地址: https://gitcode.com/gh_mirrors/cli/ClimaX

ClimaX 是微软研究院提出的天气与气候基础模型(Foundation model for weather & climate),它用大规模 CMIP6 数据预训练、再微调的方式,把「大模型」范式带进了气象领域。本文将从数据处理、参数配置到训练评估,一步步带你完整复现论文中的全球天气预报实验,即使你是第一次接触气象 AI,也能照着跑通全流程。

认识 ClimaX:天气与气候领域的「GPT」基础模型

传统气象预报模型通常为单一任务定制,而 ClimaX 走的是「先预训练、再微调」的路线:先在多个 CMIP6 气候模拟数据集上学习通用的天气演变规律,再针对具体任务微调,从而用更少的数据、更短的时间达到更强的预报效果。

在论文的全球天气预报实验中,ClimaX 以 ERA5 再分析数据为训练集,可以同时预测位势高度、温度、风场等多个变量,覆盖从 6 小时到数天不等的预报时效。复现该实验只需关注项目中两个核心文件:

  • 训练入口:src/climax/global_forecast/train.py
  • 核心配置:configs/global_forecast_climax.yaml

复现实验前的准备:环境搭建与硬件要求

实验基于 PyTorch Lightning,推荐使用 GPU 环境。先获取代码并安装依赖:

git clone https://gitcode.com/gh_mirrors/cli/ClimaX cd ClimaX conda env create --file docker/environment.yml conda activate climaX pip install -e .

💡 没有 conda 也可以参考 docs/install.md 中的 Docker 方式,两种方案都已配置好全部依赖。

硬件方面,论文实验在 8 张 GPU 上完成微调;如果你资源有限,也可以先用单卡跑通流程,再逐步扩大规模。

第一步:下载并预处理 ERA5 天气数据集

全球天气预报实验使用 WeatherBench 提供的ERA5 再分析数据(5.625° 分辨率),包含 2 米气温、10 米风场、各气压层的位势高度/温度/湿度等变量。下载完成后,需要把 NetCDF 格式转成高效的 numpy 分片文件,并顺便算出标准化所需的均值方差与气候态:

python src/data_preprocessing/nc2np_equally_era5.py \ --root_dir /mnt/data/5.625deg \ --save_dir /mnt/data/5.625deg_npz \ --start_train_year 1979 --start_val_year 2016 \ --start_test_year 2017 --end_year 2019 --num_shards 8

处理完成后会生成 train / val / test 三个子目录,以及normalize_mean.npznormalize_std.npzlat.npylon.npy等关键统计文件。数据加载的逻辑封装在 src/climax/global_forecast/datamodule.py 中,标准化和气候态都从这里读取。

第二步:看懂并修改关键训练配置

configs/global_forecast_climax.yaml 是复现的核心,你需要重点关注以下几组参数:

配置项论文默认值作用说明
data.root_dir数据集路径指向预处理好的 npz 目录
data.predict_range72预报时效(小时),可调成 6/24/72
data.out_variables5 个变量输出变量,如geopotential_5002m_temperature
model.pretrained_path预训练权重填入 5.625deg 预训练 checkpoint 路径
model.lr5e-7微调学习率,比从头训练小得多
trainer.devicesGPU 数量8 卡用 8,单卡用 1

⚠️ 论文实验的关键在于加载预训练权重:ClimaX 先在 5 个 CMIP6 数据集上完成自监督预训练,微调时通过interpolate_pos_embed自动插值位置编码(见 src/climax/utils/pos_embed.py),适配不同分辨率。

第三步:一条命令启动全球天气预报训练

修改好配置后,用一条命令启动微调训练:

python src/climax/global_forecast/train.py --config configs/global_forecast_climax.yaml \ --trainer.strategy=ddp --trainer.devices=8 \ --trainer.max_epochs=50 \ --data.root_dir=/mnt/data/5.625deg_npz \ --data.predict_range=72 \ --data.out_variables=['geopotential_500','temperature_850','2m_temperature'] \ --data.batch_size=16 \ --model.pretrained_path='你的预训练权重路径' \ --model.lr=5e-7 --model.beta_1="0.9" --model.beta_2="0.99" \ --model.weight_decay=1e-5

训练完成后,脚本会自动加载最佳 checkpoint 在测试集上评估。训练与评估的具体逻辑(含纬度加权损失、AdamW 优化器、余弦退火调度)都封装在 src/climax/global_forecast/module.py 中,值得细读。

第四步:解读评估指标与预测效果

ClimaX 使用三个气象领域标准指标评估:

  • RMSE(均方根误差):衡量预测与真值的整体偏差;
  • ACC(异常相关系数):衡量空间分布形态的吻合度,越接近 1 越好;
  • 纬度加权:由于极地网格面积小,所有指标都按纬度加权,避免高纬度过分主导误差。

指标实现见 src/climax/utils/metrics.py。训练日志会输出val/w_rmseval/w_acc等,Checkpoint 回调也会按val/w_rmse自动保存最优模型。

上图展示了 ClimaX 对 2 米气温(T2m)的预测效果:初始条件、6 小时真实值、6 小时预测值与偏差四栏对比,可以看到预测与真实演变高度一致,偏差区域也控制在很小范围。更多预测可视化可查看docs/assets/images/gifs/目录下的z500.gift850.gif等动图。

常见问题与调优技巧

  • 训练不收敛?检查pretrained_path是否加载成功,微调时学习率应比预训练低一个量级(5e-7 左右)。
  • 显存不足?调小batch_size,或把precision保持为 16 混合精度。
  • 想试更长预报时效?predict_range从 72 加大到 120/240,观察 RMSE 随时效的增长曲线,这正是论文中「预报时效越长误差越大」的经典结论。
  • 想从头训练?--model.pretrained_path=""置空即可,但需要更长的训练时间。

结语

至此,你已经完整走通了 ClimaX 全球天气预报实验的五个环节:数据预处理 → 配置修改 → 微调训练 → 指标评估 → 可视化验证。这套「预训练 + 微调」范式还可以继续延伸到区域预报与气候预估场景,相关代码分别位于 src/climax/regional_forecast/ 与 src/climax/climate_projection/。上手跑通一次实验,你就真正掌握了这款天气与气候基础模型的核心玩法。🚀

【免费下载链接】ClimaXFoundation model for weather & climate项目地址: https://gitcode.com/gh_mirrors/cli/ClimaX

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

← 返回列表