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

2023-12-28 19:29:45

今年涌现了诸如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、其他问题?欢迎交流!如有不当之处,欢迎指正!

文章来源:https://blog.csdn.net/yutingzhaomeng/article/details/135269301
本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。