今年涌现了诸如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.10
source activate large
2、按照https://github.com/ecmwf-lab/ai-models安装所需库
pip install ai-models
3、添加EC封装好的各种大模型,目前支持的包括GraphCast、FourCastNet、Pangu等,直接pip安装即可
pip install ai-models-panguweather
pip install ai-models-fourcastnet
pip install ai-models-graphcast # Install details at https://github.com/ecmwf-lab/ai-models-graphcast
pip 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.git
cd ai-models-fuxi
python 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.whl
pip 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即上图里的
fourcastnet
fourcastnetv2-small
fuxi
graphcast
panguweather
也可以从各家大模型在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/v2
key:?对应的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、其他问题?欢迎交流!如有不当之处,欢迎指正!