如何在本地部署现有气象大模型

今年涌现了诸如Pangu、Fuxi、Fengwu、GraphCast、FourCastNet等诸多气象大模型，本文介绍如何用EC开发的ai-models在本地部署以上模型。

本文测试环境系统为：

Ubuntu 18.04.6 LTS + Anaconda 3 + Cuda 11.8 + libcudnn 8

1、创建并启动虚拟环境

conda create -n large python=3.10source activate large

2、按照https://github.com/ecmwf-lab/ai-models安装所需库

pip install ai-models

3、添加EC封装好的各种大模型，目前支持的包括GraphCast、FourCastNet、Pangu等，直接pip安装即可

pip install ai-models-panguweatherpip install ai-models-fourcastnetpip install ai-models-graphcast  # Install details at https://github.com/ecmwf-lab/ai-models-graphcastpip install ai-models-fourcastnetv2

注：如果GraphCast用pip一键安装不成功（大概率是），可以按照https://github.com/ecmwf-lab/ai-models-graphcast提到的方式，首先安装jax，再手动安装GraphCast

pip install -r requirements-gpu.txt -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

4、添加Fuxi大模型，当前ai-models支持Fuxi，但貌似还不支持pip一键安装，可以按如下方式手动添加

git clone https://github.com/tpys/ai-models-fuxi.gitcd ai-models-fuxipython setup.py install

5、添加Fengwu大模型，Fengwu目前还未对ai-models做适配，这个需要手动添加。可以按照EC或Fuxi的风格修改一下，也可以直接从如下链接下载编译好的whl文件

链接: https://pan.baidu.com/s/1WCwJ4ptLXHeNHHO3xw_L4A?pwd=tzyu 提取码: tzyu 复制这段内容后打开百度网盘手机App，操作更方便哦链接: https://pan.baidu.com/s/1sOFN501Q-pKbUBiowZU0Ng?pwd=tzyu 提取码: tzyu 复制这段内容后打开百度网盘手机App，操作更方便哦

下载好whl后，直接pip 安装即可

pip install ai_models_fengwu-0.0.1-py3-none-any.whlpip install ai_models_fengwuv2-0.0.1-py3-none-any.whl

上述两个whl分别对应Fengwu模型中提到的without Transfer和with Transfer

6、至此，ai-models支持的大模型全部安装完成，可以通过如下命令查看是否安装成功

ai-models --models

如全部安装成功，则出现如下结果

如有未安装成功的（即未显示的，考虑重新安装一下）

7、关于预训练模型，ai-models对FourCastNet、GraphCast、Pangu、Fuxi均支持自动下载，按如下方式即可

ai-models --download-assets <model_name>

model_name即上图里的

fourcastnetfourcastnetv2-smallfuxigraphcastpanguweather

也可以从各家大模型在github公开的链接手动下载，如Fuxi，以下链接可以下载

https://github.com/tpys/ai-models-fuxi

而对Fengwu，目前需要手动下载，下载链接

https://github.com/OpenEarthLab/FengWu

8、接入输入数据。目前测试了通过CDS访问ERA5的方式，具体来说，首先到如下链接注册账号，

https://cds.climate.copernicus.eu/

然后点击自己的用户名，找到如下

在本地用户根目录下创建文件.cdsapirc，将

url: https://cds.climate.copernicus.eu/api/v2key: 对应的API Key

添加并保存即可。

9、如何运行？如采用上述步骤8提到的CDS数据输入，按如下方式运行即可

ai-models --input cds --date 20230110 --time 0000 fuxi --assets fuxi

其中，input指定了数据输入方式为CDS，程序会自动下载date和time指定时次的文件，并保存至本地。fuxi即模型名，也可替换为上述步骤6提到的任意模型，assets为下载预训练模型对应的目录。

其他如指定输入文件的运行方式待后续更新。

默认做未来10天240小时逐6小时的预报，如无需太多预报时次，可通过lead_time参数传递指定，更多参数可通过ai-models --help命令获取。

10、效果如何？见下图

单从模型结构的角度看，个人还是特别喜欢GraphCast的设计。

11、关于安装过程可能遇到的问题：

1、pangu、fuxi、fengwu等模型无法调用gpu解决方案：注意看提示的问题，一般是cuda、libcudnn版本不对应，按本文开始提到的环境配置亲测可用。此外，如采用CPU，单模型运行时间一般为8小时左右不等，采用单卡GPU，一般不超过5分钟。​2、GraphCast提示ptx问题解决方案：也是版本问题​3、单模型运行一般占用10G左右的显存，对pangu，因为有6和24两个模型，可能无法实现单卡运行解决方案：用空间换算力，将6和24两个模型分别指定至两个GPU，需要修改源码，找到~/anaconda3/envs/large/lib/python3.10/site-packages/ai-models-panguweather/model.py文件，第81、88行，分别增加 ort_session_*.set_providers(['CUDAExecutionProvider'],  provider_options=[{'device_id': *}])​4、其他问题？欢迎交流！如有不当之处，欢迎指正！