手把手教你创建自己的人工智能天气预报！

在本文中，我们展示了如何使用ECMWF提供的AI-Models存储库来生成AI天气预报。使用来自气候数据存储的初始化数据，我们能够使用两个模型生成预报：华为的PanguWeather和NVIDIA的FourCastNet。使用预训练模型生成一个为期10天的天气预报在GPU上大约需要2分钟，在CPU上大约需要3小时。与生成数值天气预报所需的巨大计算工作相比，推理的成本非常低。现在，个人消费者和资金较

七七Seven～

2183人浏览 · 2024-04-14 10:39:46

七七Seven～ · 2024-04-14 10:39:46 发布

手把手教你创建自己的人工智能天气预报

基于数据驱动的天气预报使用预训练模型，成本低廉且能够提供与传统数值天气模式相当准确的预报。多家公司和研究实验室已经开发了人工智能天气模型，包括：

PanguWeather [华为]

FourCastNet [英伟达]

GraphCast [谷歌DeepMind]

欧洲中期天气预报中心（ECMWF）提供了使用这些模型生成天气预报的例程[1]。即使在笔记本电脑上也可以执行模型推理，但建议使用GPU。

在这篇文章中，目的如下：

展示给你如何创建自己的人工智能天气预报
比较在GPU和CPU上的模型推理时间
可视化PanguWeather和FourCastNet的天气预报，包括温度、水汽和急流

1 背景信息

传统上，天气预报依赖于在全球网格上解决的数值天气模型。这需要大量的计算资源，世界上只有少数几个天气服务机构具备生成全球天气预报的能力。

AI天气模型是基于过去的天气数据进行训练的，称为再分析数据。利用训练好的模型和当前的天气状态，可以仅使用数值天气预报所需计算资源的一小部分来生成天气预报。

该示意图展示了人工智能天气预报如何利用当前天气作为输入来预测未来天气:

AI天气预报示意图。左：2021年01月01日00:00 UTC的风场。右：2021年01月01日23:00 UTC的风场。数据来源：Copernicus气候数据存储的ERA5 [3]

本文涵盖的所有三个人工智能天气模型都提供0.25°（25公里）的空间分辨率和6小时的时间分辨率的天气预报。

2 开始

2.1 代码和环境

从ECMWF的Github页面[1]获取代码以运行预训练的天气模型，并按照安装说明进行操作。我在本教程中使用的是v0.2.5版本。为所需的Python包设置一个 conda 环境。

conda create -n ai-models python=3.10  
conda activate ai-models  
conda install cudatoolkit  
pip install ai-models

使用可用的AI模型进行安装

pip install ai-models-panguweather  
pip install ai-models-fourcastnet

GraphCast需要一个特殊的安装程序，下面会进行介绍。

2.2 预训练模型权重

需要下载PanguWeather、GraphCast和FourCastNet的预训练模型权重

ai-models --download-assets --assets assets-panguweather panguweather  
ai-models --download-assets --assets assets-fourcastnet  fourcastnet  
ai-models --download-assets --assets assets-graphcast    graphcast

2.3 初始化数据

AI天气模型需要使用当前的天气数据进行初始化，以便进行预测。获取数据有两种方式，一种是从ECMWF服务MARS获取，另一种是从Copernicus气候数据存储库（CDS）获取。两项服务都免费提供初始化数据。

正在使用CDS的再分析数据进行本教程。按照他们网站上的说明获取API密钥以访问数据[2]。请注意，再分析数据会有大约五天的延迟。

2.4 创建人工智能天气预报

我们现在已经准备好生成我们自己的人工智能天气预报了。唯一剩下的就是选择一个时间和日期进行初始化，并使用预训练模型开始模型推理。

2.4.1 PanguWeather

要求基于CDS初始化数据的PanguWeather预报，从2023年09月20日00:00 UTC开始，请使用以下方法：

ai-models --input cds --date 20230920 --time 0000 --assets assets-panguweather panguweather

这些日志非常全面。预训练模型 pangu_weather_6.onnx 可以提前6小时生成天气预报，这里将其迭代应用40次，总共达到240小时（10天）。

2023-09-26 12:06:45,724 INFO Loading assets-panguweather/pangu_weather_24.onnx: 26 seconds.  
2023-09-26 12:07:10,289 INFO Loading assets-panguweather/pangu_weather_6.onnx: 24 seconds.  
2023-09-26 12:07:10,289 INFO Model initialisation: 52 seconds  
2023-09-26 12:07:10,290 INFO Starting inference for 40 steps (240h).

对于每次迭代，日志会告诉用户花费了多少时间。在实验中，ONNX运行时无法找到GPU，因此推理是在CPU上进行的，每个天气预测步骤大约需要3分15秒。

2023-09-26 12:10:27,260 INFO Done 1 out of 40 in 3 minutes 16 seconds (6h), ETA: 2 hours 11 minutes 18 seconds.

2.4.2 FourCastNet

为了在相同的初始化时间生成FourCastNet的预测，同时将训练好的模型权重下载到子目录 assets-fourcastnet 中，我们使用：

ai-models --download-assets --assets assets-fourcastnet --input cds --date 20230920 --time 0000 fourcastnet

再次，日志为我们提供了关于模型推断的全面信息。现在，ONNX运行时能够找到GPU，并且推断速度更快，只需要2分钟生成10天的预测。

2023-09-26 13:23:48,633 INFO Using device 'CUDA'. The speed of inference depends greatly on the device.  
2023-09-26 13:23:59,710 INFO Loading assets-fourcastnet/backbone.ckpt: 13 seconds.  
2023-09-26 13:24:10,673 INFO Loading assets-fourcastnet/precip.ckpt: 10 seconds.  
2023-09-26 13:24:10,733 INFO Model initialisation: 46 seconds  
2023-09-26 13:24:10,733 INFO Starting inference for 40 steps (240h).  
2023-09-26 13:24:14,247 INFO Done 1 out of 40 in 3 seconds (6h), ETA: 2 minutes 20 seconds.

2.4.3 GraphCast

对于GraphCast，我们首先需要按照https://github.com/ecmwf-lab/ai-models-graphcast中的安装步骤进行安装。然后，我们可以使用以下命令生成天气预报：

ai-models --download-assets --assets assets-graphcast --input cds --date 20230920 --time 0000 graphcast

由于一个在GitHub上仍然存在的错误（在[1]中的问题＃10），暂时无法使用GraphCast完成预测。

2.5 预测输出文件

AI天气模型生成的输出文件采用GRIB格式。对于240小时的预报，文件大小为：

2.3 GB的FourCastNet（3个气压层：850 hPa，500 hPa，250 hPa）
5.4 GB的PanguWeather（13个气压层）

两个模型都输出关键的气象变量：

温度
风速和风向

PanguWeather还提供位势高度和湿度，而FourCastNet提供水汽含量。

2.6 生成图表

这些由人工智能生成的预测的适当评估和基准测试是一个复杂的话题，超出了本文的范围。在这里，我们重点关注两个人工智能天气模型之间的视觉比较。

为了对天气预报进行视觉分析，我们选择温度变量和初始化后 24 小时的预报。这是生成图表的模板：

import xarray as xr  
import seaborn as sns  
from matplotlib import pyplot as plt  
import cartopy.crs as ccrs  
  
# load data  
ds_fourcast = xr.open_dataset('fourcastnet.grib')  
  
ix_step = 3  # 0 = first step in weather prediction  
ix_level = 0 # 0 = surface pressure level  
temp_degC = ds_fourcast['t'][ix_step, ix_level] - 273.15 # subtract 273.15 to convert from Kelvin to Celsius.  
  
# plot settings  
sns.set_style('ticks')  
tmin = -45  
tmax = +45  
nlev = 19  
  
  
# generate figure  
fig = plt.figure()  
ax = fig.add_subplot(1, 1, 1, projection=ccrs.PlateCarree())  
ax.set_global()  
ax.coastlines('110m', alpha=0.1)  
  
# contour plot  
img = ax.contourf(ds_fourcast.longitude.values, ds_fourcast.latitude.values,   
                  temp_degC.values,   
                  levels = np.linspace(tmin, tmax, nlev),  
                  cmap='RdBu_r',  
                  transform=ccrs.PlateCarree()  
                 )  
  
# hours passed since initialization  
dt_hours = ((ds_fourcast['valid_time'][ix_step] - ds_fourcast['time']).values / 1e9 / 3600).astype(float)  
  
# annotation text  
ax.text(-180, 90,   
        f'After {dt_hours} hours ({ds_fourcast["valid_time"][ix_step].values})',   
        transform=ccrs.PlateCarree(),  
        va='bottom', ha='left',  
       )  
  
plt.colorbar(img, orientation='horizontal', label='850 hPa temperature (°C)')  
plt.show()

2.6.1 温度

以下图表比较了PanguWeather和FourCastNet在850 hPa层次上的温度。这个层次并不直接位于地表，但它是两个输出文件中唯一包含的层次。从视觉上看，我们可以看到温度场显示出相似的模态。

850 hPa 24小时后的温度。

2.6.2 水汽

FourCastNet提供了大气水汽总量，可以与携带雨滴的云量进行比较。

在下图中，我们比较了这个数量在四个不同的快照中的情况：1天后（24小时），3天后（72小时），5天后（96小时），以及10天后（240小时）。

FourCastNet预测后1、3、5和10天的大气水汽总量。

我们观察到在深度学习中一个众所周知的现象：由于模型是使用均方根误差（RMSE）进行训练的，它有产生平滑场的倾向。随着时间的推移，当模型被迭代应用时，预测变得不那么明显。

10天后，我们甚至在南美洲的南端观察到了人工制品。请注意，AI天气模型目前不使用生成式AI技术，因为它们更注重物理准确性而非视觉吸引力。

2.6.3 风速

PanguWeather提供不同气压层次下的风速。我们选择250 hPa的压力水平，该水平对应着急流最为明显的层次。

急流是一种环极气流，影响着美国和欧洲等中纬度地区的天气。图表显示了以南半球急流为重点的风速情况，在我预测的时间段内，急流更加明显。

南半球急流预报与盘古天气。

为了创建正交投影并计算风速，我们使用了以下模板：

import xarray as xr  
import seaborn as sns  
from matplotlib import pyplot as plt  
import cartopy.crs as ccrs  
  
# load data  
#ds_pangu = xr.open_dataset('/home/k/k202141/rootgit/ai-models/panguweather.grib')  
  
ix_step = 3  # 0 = first step in weather prediction  
ix_level = 8 # 0 = surface pressure level, 2 - 850 hPa level  
windspeed = (ds_pangu['u'][ix_step, ix_level]**2 + ds_pangu['u'][ix_step, ix_level]**2)**0.5  
  
# plot settings  
sns.set_style('ticks')  
tmin = 0  
tmax = 150  
nlev = 16  
  
# generate figure  
fig = plt.figure()  
ax = fig.add_subplot(1, 1, 1, projection=ccrs.Orthographic(central_latitude=-30))  
ax.set_global()  
ax.coastlines('110m', alpha=1.0)  
  
# contour plot  
img = ax.contourf(ds_pangu.longitude.values, ds_pangu.latitude.values,   
                  windspeed.values,   
                  levels = np.linspace(tmin, tmax, nlev),  
                  cmap='gist_ncar',  
                  transform=ccrs.PlateCarree()  
                 )  
  
# hours passed since initialization  
dt_hours = ((ds_pangu['valid_time'][ix_step] - ds_pangu['time']).values / 1e9 / 3600).astype(float)  
print(dt_hours)  
  
plt.colorbar(img, orientation='horizontal', label='Wind speed (m/s)', shrink=0.5)  
plt.show()