全球天气预报模型Fuxi运行debug经验分享

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.11
torch = 2.1.0
torch_npu = 2.1.0 
timm = 0.9.10
numpy = 1.23.5

5.3 在WeatherLearn目录下 新建fuxi.py

内容为

from weatherlearn.models import Fuxi
import torch
import torch_npu
if __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 = True
torch.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占用