基于昇腾适配数据驱动的全球天气预报模型 Fuxi

1. 摘要

Fuxi 是由复旦大学的研究人员开发的一个基于数据驱动的全球天气预报模型，主要由 Cube Embedding、U-Transformer 和全连接层构成。Fuxi 摒弃了传统复杂的微分方程，转而通过多阶段机器学习架构，可提供 15 天的全球预报。时间分辨率为 6 小时，空间分辨率为 0.25°。

2. Fuxi 介绍

Fuxi 模型是由复旦大学的研究人员开发的一个基于数据驱动的全球天气预报模型，它摒弃了传统复杂的微分方程，转而通过多阶段机器学习架构，可提供 15 天的全球预报。时间分辨率为 6 小时，空间分辨率为 0.25°，相当于赤道附近约 25 公里 x 25 公里的范围，使用 ECMWF39 年的 ERA5 再分析数据集训练，在 15 天预报尺度上实现了效率与精度的双重突破。

3. 网络架构

基本的伏羲模型体系结构由三个主要组件组成，如图所示：Cube Embedding、U-Transformer 和全连接层。输入数据结合了上层空气和地表变量，并创建了一个维度为 69×720×1440 的数据立方体，以一个时间步作为一个 step。高维输入数据通过联合时空 Cube Embedding 进行维度缩减，转换为 C×180×360。Cube Embedding 的主要目的是减少输入数据的时空维度，减少冗余信息。随后，U-Transformer 处理嵌入数据，并使用简单的全连接层进行预测，输出首先被重塑为 69×720×1440。

• Cube Embedding 为了减少输入数据的空间和时间维度，并加快训练过程，应用了 Cube Embedding 方法。具体地，空时立方体嵌入采用了一个三维（3D）卷积层，卷积核和步幅分别为 2×4×4，输出通道数为 C。在空时立方体嵌入之后，采用了层归一化（LayerNorm）来提高训练的稳定性。最终得到的数据立方体的维度是 C×180×360。

• U-Transformer U-Transformer 还包括 U-Net 模型的下采样和上采样块。下采样块在图中称为 Down Block，将数据维度减少为 C×90×180，从而最小化自注意力计算的计算和内存需求。Down Block 由一个步长为 2 的 3×3 2D 卷积层和一个残差块组成，该残差块有两个 3×3 卷积层，后面跟随一个组归一化（GN）层和一个 Sigmoid 加权激活函数(SiLU)。SiLU 加权激活函数通过将 Sigmoid 函数与其输入相乘来计算。 上采样块在图中称为 Up Block，它与 Down Block 使用相同的残差块，同时还包括一个 2D 反卷积，内核为 2，步长为 2。Up Block 将数据大小缩放回 C×180×360。此外，在馈送到 Up Block 之前，还包括一个跳跃连接，将 Down Block 的输出与 Transformer Block 的输出连接起来。 中间结构是由 18 个重复的 Swin Transformer 块构建而成，通过使用残差后归一化代替前归一化，缩放余弦注意力代替原始点积自注意力，Swin Transformer 解决了诸如训练不稳定等训练和应用大规模的 Swin Transformer 模型会出现几个问题。

4. 核心设计

• 全局粗预测：首阶段模型基于历史气象数据（气压、温度、风速）,生成未来 15 天的低分辨率趋势框架。

• 局部细修正：第二阶段针对台风、暴雨等高影响天气，在关键区域进行空间精细化调整，分辨率提升至 0.25°网格。

• 动态反馈：最终阶段引入实时观测数据流，持续校准预测轨迹，抑制长期预报中的误差扩散。

这种级联架构使得 Fuxi 在保持 AI 模型高效推理的同时（单次预测仅需数分钟），关键指标如 500 hPa 位势高度异常相关系数（ACC）达到 0.82，性能显著优于世界上最佳的基因物理的天气预报系统——欧洲中期天气预报中心（ECMWF）的高分辨率预报系统。其训练过程中嵌入了质量守恒、能量平衡等物理规则约束，确保预测结果既符合数据规律，又不违背气象学基本原理。作为首个支持 15 天全球高分辨率预报的开源 AI 模型，Fuxi 不仅为气象业务提供了快速响应的新工具，其模块化设计也为极端天气预警、航空路径优化等场景提供了可定制化基础。

5. 实验

5.1 下载源码

https://github.com/lizhuoq/WeatherLearn

5.2 安装依赖库

python = 3.11torch = 2.1.0torch_npu = 2.1.0 timm = 0.9.10numpy = 1.23.5

5.3 在 WeatherLearn 目录下 新建 fuxi.py

内容为

from weatherlearn.models import Fuxiimport torchimport torch_npuif __name__ == '__main__':    device = torch.device('npu' if torch_npu.npu.is_available() else 'cpu')    B = 1  # batch_size    in_chans = out_chans = 70  # number of input channels or output channels    input = torch.randn(B, in_chans, 2, 721, 1440).to(device)  # B C T Lat Lon    fuxi = Fuxi().to(device)    output = fuxi(input)  # B C Lat Lon    print(output.shape)Å

5.4 source cann

source /usr/local/Ascend/ascend-toolkit/set_env.sh

5.5 运行脚本

python3 fuxi.py

报错 Conv3dv2 only support static shape

解决办法

在实例化模型前添加

torch.npu.config.allow_internal_format = Truetorch.npu.set_compile_mode(jit_compile=True）

再次运行，报错 out of memory

解决办法 1

禁用梯度计算和降低计算精度（混合精度）

vi fuxi.py

在第 21、22 行添加

with torch.no_grad():    with torch.npu.amp.autocast():

解决办法 2

减少网络结构的层数

vi /WeatherLearn/weatherlearn/models/fuxi/fuxi.py

在 144 行 Fuxin 类的初始化中添加参数 depth=12，将 161 行的 depth=48 修改为 depth

再次运行

npu 占用