HFAiLab/OpenCastKit

The open-source solutions of FourCastNet and GraphCast

OpenCastKit: an open-source solutions of global data-driven high-resolution weather forecasting

简体中文 | English

本项目是由幻方AI团队复现优化，并开源的全球AI气象预报模型工具库。基于 FourCastNet 和 GraphCast 的论文，我们构建了一个新的全球AI气象预报项目——OpenCastKit，它能够与欧洲中期天气预报中心（ECMWF）的传统物理模型——高分辨率综合预测系统（IFS），进行直接比较。

我们将基于1979年1月到2022年12月的ERA5数据训练出来的模型参数开源到 Hugging Face 仓库中，并上线了一个每日更新的 HF-Earth，展示模型的预测效果。

下面是一些预测案例：

依赖

• hfai >= 7.9.5
• torch >=1.8

训练

原始数据来自欧洲中期天气预报中心（ECMWF）提供的一个公开可用的综合数据集 ERA5 ，需要通过 data_factory/convert_ear5_hourly.py 脚本提出数据特征，转化为高性能训练样本格式 FFRecord 下的样本数据。

训练 FourCastNet

本地运行：

   python train_fourcastnet.py --pretrain-epochs 100 --fintune-epochs 40 --batch-size 4

也可以提交任务至幻方萤火集群，使用96张A100进行数据并行训练

   hfai python train_fourcastnet.py --pretrain-epochs 100 --fintune-epochs 40 --batch-size 4 -- -n 12 --name train_fourcastnet

训练 GraphCast

本地运行：

   python train_graphcast.py --epochs 200 --batch-size 2

也可以提交任务至幻方萤火集群，使用256张A100进行流水线并行训练

   hfai python train_graphcast.py --epochs 200 --batch-size 2 -- -n 32 --name train_graphcast

引用

@article{pathak2022fourcastnet,
  title={Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators},
  author={Pathak, Jaideep and Subramanian, Shashank and Harrington, Peter and Raja, Sanjeev and Chattopadhyay, Ashesh and Mardani, Morteza and Kurth, Thorsten and Hall, David and Li, Zongyi and Azizzadenesheli, Kamyar and others},
  journal={arXiv preprint arXiv:2202.11214},
  year={2022}
}

@article{remi2022graphcast,
  title={GraphCast: Learning skillful medium-range global weather forecasting},
  author={Remi Lam, Alvaro Sanchez-Gonzalez, Matthew Willson, Peter Wirnsberger, Meire Fortunato, Alexander Pritzel, Suman Ravuri, Timo Ewalds, Ferran Alet, Zach Eaton-Rosen, Weihua Hu, Alexander Merose, Stephan Hoyer, George Holland, Jacklynn Stott, Oriol Vinyals, Shakir Mohamed, Peter Battaglia},
  journal={arXiv preprint arXiv:2212.12794},
  year={2022}
}