问海 WenHai 模型复现流程

本文记录一次问海（WenHai）全球海洋预报模型的完整复现。目标是从公开再分析资料出发，构造模型要求的海洋初始场和大气强迫，在 GPU 服务器上运行官方 ONNX 权重，把一次性的示例脚本整理成能够覆盖 2023 年、可校验、可断点续跑的工程流程。复现聚焦于用官方权重做推理。

问海可以看作一个混合 AI 模型：Swin Transformer 负责学习全球海洋状态的演变，海气界面的热量、动量和淡水通量则仍由 bulk formulae 显式计算。它把物理公式算出的海气通量作为外部条件，结合神经网络学到的海洋状态演变规律，逐日滚动预报海洋状态。

工程内容与复现目标

当前工程的核心文件如下：

文件或目录	作用
WenHai.onnx	约 251 MB 的官方 ONNX 模型权重
inference.py	官方单次起报推理脚本
min_GLORYS.npy、max_GLORYS.npy	93 通道海洋状态的归一化边界
min_flux.npy、max_flux.npy	8 通道海气通量的归一化边界
mask_GLORYS.nc	与 GLORYS 网格一致的逐通道海陆掩膜
data_prep/	下载、重网格、批次规划和输入校验
server/	AutoDL 环境安装、批量推理、压缩输出和断点续跑
ncomm_sourcedata/	论文图表对应的源数据

官方示例只演示“一份初始场 + 一份 10 天强迫”的运行方式。本次复现进一步做了三件事：

1. 自动准备 2023 年所需的 GLORYS12 和 ERA5 输入；
2. 每 10 天起报一次、每次预报 10 天，用 37 个批次连续覆盖全年；
3. 把长时间任务改造成可检查、可压缩、可在单个预报日处恢复的批处理流程。

这里的“复现”特指使用作者发布的模型权重完成推理，并不包含训练数据重建和网络重新训练。

问海究竟吃进去什么

93 通道海洋状态

海洋初始场来自 Copernicus Marine 的 GLORYS12 全球海洋再分析，水平分辨率为 $1/{12}^{\circ }$，网格大小为 $2041\times 4320$。模型使用五个变量：

变量	含义	垂向层数
uo	纬向流速	23
vo	经向流速	23
thetao	海水位温	23
so	盐度	23
zos	海表高度	1

四个三维变量与一个二维变量沿通道维拼接后得到：

$$23+23+23+23+1=93\text{channels}.$$

因此单个时刻的模型状态形状为：

[batch, channel, latitude, longitude]
= [1, 93, 2041, 4320]

GLORYS 原始深度轴不能直接照目录顺序切片。官方给出的 23 层索引按“表层为第 1 层”定义，而数据目录可能按深到浅返回。预处理脚本先执行 sortby("depth")，再选择指定索引，并与 0.494—643.567 m 的已知深度表核对。这个检查非常重要：深度方向即使完全反了，数组形状仍然可能正确，但模型会把深海当成表层。

下面是仓库中 mask_GLORYS.nc 的表层掩膜。原文件是与 93 个状态通道逐一对应的布尔数组，逻辑尺寸为 [93, 2041, 4320]。

8 个大气变量与 8 个通量通道

大气强迫来自 ERA5，使用 8 个近地面变量：

ERA5 变量	含义
u10、v10	10 m 纬向和经向风
t2m、d2m	2 m 气温和露点温度
msl	平均海平面气压
ssr	地表净短波辐射
strd	地表向下长波辐射
mtpr	平均总降水率

脚本先把逐小时 ERA5 求为逐日平均，再用双线性插值映射到 GLORYS 的 $1/{12}^{\circ }$ 网格。经度方向额外做周期补边，避免在 ${180}^{\circ }$ 日界线处留下 NaN 接缝。

ERA5 经过转换后才进入模型：推理时会结合当前的表层流速和海表温度，通过 aerobulk 的 NCAR bulk formulae 计算长波净辐射、短波辐射、感热、潜热、纬向风应力、经向风应力、蒸发和降水共 8 个通量通道，这 8 个通量通道才是模型真正的第二个输入：

$$F_t=[q_l,q_{sw},q_h,q_e,{\tau }_x,{\tau }_y,E,P].$$

其中比湿由露点温度和海平面气压计算，长波净辐射使用

$$q_l=\mathrm{strd}-\sigma T_s^4,$$

然后所有通量按预先保存的极值归一化到 $[0,1]$。

残差式自回归预报

ONNX 网络输出的是归一化空间中的状态增量。推理一步可以概括为：

$${\tilde{X}}_{t+1}={\mathrm{clip}}_{[0,1]}[M\odot ({\tilde{X}}_t+f_{\theta }({\tilde{X}}_t,{\tilde{F}}_t))].$$

式中 ${\tilde{X}}_t$ 是归一化后的 93 通道状态，${\tilde{F}}_t$ 是 8 通道通量，$M$ 是海陆掩膜。得到第 1 天结果后，它会成为第 2 天的输入；如此重复，形成 10 天自回归预报。

这也解释了一个容易弄错的时间约定：若初始场日期为 2023-01-01，则强迫文件必须包含 01-01—01-10，输出日期才是 01-02—01-11。强迫的起始日与初始场当天相同，早于预报输出的第一天。

2023 年批量预报方案

本次复现设置每 10 天起报一次，每次预报 10 天：

第 1 次：初始场 2023-01-01 -> 预报 01-02 ... 01-11
第 2 次：初始场 2023-01-11 -> 预报 01-12 ... 01-21
...
第 37 次：初始场 2023-12-27 -> 预报 12-28 ... 2024-01-06

起报步长等于预报长度，所以相邻批次无缝且不重叠。默认配置连续覆盖 2023-01-02—2024-01-06；2023-01-01 本身没有预报值。若必须补齐这一天，可把 INCLUDE_SPINUP_INIT 设为 True，额外加入 2022-12-22 起报的一次 spin-up。

数据量和硬件预算

这一步不能只看 251 MB 的权重文件。真正占空间的是全球 $1/{12}^{\circ }$ 的输入和逐日输出：

项目	估算规模
37 份 GLORYS 初始场	约 40 GB（压缩后）
37 份 10 日 ERA5 强迫	约 53 GB（压缩后）
370 个逐日预报，未压缩	约 1.21 TB
370 个逐日预报，deflate-1	约 600—700 GB

单个预报日的未压缩状态约为：

$$93\times 2041\times 4320\times 4\mathrm{B}\approx 3.28\mathrm{GB}.$$

推理期间还要同时保存归一化状态、掩膜、模型输出和写盘副本，内存峰值预计约 25—35 GB；显存除 1.6 GB 左右的 FP16 输入外，还要容纳 Swin Transformer 的中间激活。实际开跑前应先用一个预报日做冒烟测试，确认无误后再提交全年任务。

先在 Windows 端准备好输入数据

数据准备和 GPU 推理使用两套独立环境。下载侧依赖主要来自 PyPI，适合用 uv；推理侧的 aerobulk-python 包含 Fortran 扩展，使用 conda-forge 更省事。

进入数据准备目录并安装依赖：

cd D:\Documents\WenHai-1.0\data_prep
uv sync

ERA5 主路径使用公开、匿名可读的 ARCO-ERA5 v3 Zarr；GLORYS12 必须使用免费的 Copernicus Marine 账号：

uv run copernicusmarine login

设置数据盘后先只查看计划：

$env:WENHAI_DATA_ROOT = 'E:\WenHai_data'
uv run python prepare_2023.py --dry-run

确认 37 个起报日期、预计容量和目标目录后，再开始下载：

uv run python prepare_2023.py

脚本支持灵活补跑：

uv run python prepare_2023.py --runs 1-6
uv run python prepare_2023.py --runs 1,5,9
uv run python prepare_2023.py --only era5
uv run python prepare_2023.py --force

每个成品先写入 .nc.part，结束后再原子替换成 .nc。重新运行时，已存在且通过形状检查的文件会被跳过，所以网络中断后可以直接续跑。

最终目录结构为：

E:\WenHai_data\
├─ glorys\GLORYS_23lev_20230101.nc
├─ era5\ERA5_d0.083_20230101_10d.nc
├─ logs\prepare_YYYYmmdd_HHMMSS.log
├─ manifest_2023.json
└─ run_inference_2023.txt

先验证数据，再动用 GPU

仅检查维度远远不够。verify.py 还会检查：

• 8 个 ERA5 变量的数值范围和有限值；
• 经度首尾列是否连续，防止日界线插值错误；
• GLORYS 的 23 个深度是否表层在前；
• 温度、盐度、流速、海表高度的合理范围和湿点比例；
• 全球平均温度是否随深度总体降低。

uv run python verify.py
uv run python verify.py --init 20230101

形状正确但深度翻转、单位漂移或日界线出现 NaN 的文件，都可能在数小时推理后才以“结果完全不合理”的形式暴露。把数值校验放在上传和推理之前，代价远小于重新运行。

在 AutoDL 上配置推理环境

将仓库、模型辅助文件和准备好的数据解压到本地数据盘，逐日读写全部在本地进行：

cd /root/autodl-tmp
unzip -q /root/autodl-fs/WenHai/WenHai_data.zip -d /root/autodl-tmp/
unzip -q /root/autodl-fs/WenHai/WenHai-1.0.zip -d /root/autodl-tmp/

cd /root/autodl-tmp/WenHai-1.0
cp /root/autodl-fs/WenHai/{WenHai.onnx,mask_GLORYS.nc,*_GLORYS.npy,*_flux.npy} .
df -h /root/autodl-tmp

安装推理环境：

cd /root/autodl-tmp/WenHai-1.0/server
bash setup_env.sh
conda activate wenhai

setup_env.sh 从 conda-forge 安装 aerobulk-python、xarray、NetCDF4、MetPy 等依赖，再安装 onnxruntime-gpu 及匹配的 CUDA/cuDNN wheel，并在环境激活时补充 LD_LIBRARY_PATH。自检会确认 bulk formulae 能返回通量，以及 ONNX Runtime 能看到 CUDAExecutionProvider。

冒烟测试通过后再跑全年推理

先做只读规划：

cd /root/autodl-tmp/WenHai-1.0/server
python run_forecast.py --dry-run

然后只运行第 1 个起报的第 1 天：

python run_forecast.py --runs 1 --max-lead 1

另开终端观察 nvidia-smi -l 2 和 free -g。单日跑通后，再用 xarray 检查输出不是全 NaN 或归一化边界值：

import xarray as xr

path = (
    "/root/autodl-tmp/WenHai_data/forecast/"
    "init_20230101_10day/fcst20230102/"
    "fcst20230102_lead1_byWenHai.nc"
)
ds = xr.open_dataset(path)
sst = ds.thetao.isel(time=0, depth=0)
print("SST min/mean/max:", float(sst.min()), float(sst.mean()), float(sst.max()))
print("SSH min/mean/max:", float(ds.zos.min()), float(ds.zos.mean()), float(ds.zos.max()))

表层温度通常应处在约 $-2$—$35{}^{\circ }\mathrm{C}$，海表高度应为约 $\pm 2$ m 的量级。若全是 NaN，或几乎全部卡在 0/1 附近，应优先检查掩膜与归一化文件是否齐全。

确认资源和数值都正常后，在 tmux 中启动全年任务：

tmux new -s wenhai
conda activate wenhai
cd /root/autodl-tmp/WenHai-1.0/server
python run_forecast.py 2>&1 | tee /root/autodl-tmp/WenHai_data/logs/inference.log

也可以分批执行：

python run_forecast.py --runs 1-6
python run_forecast.py --runs 7,8,20

为什么没有直接循环调用官方脚本

官方 inference.py 很适合验证单个样例，但直接循环 37 次会放大几个工程问题。server/run_forecast.py 保持数值计算路径不变，主要调整了任务管理：

改动	作用
ONNX 会话和进程池只创建一次	避免 37 次重复加载模型、370 次重复创建 worker
按预报日识别最长连续前缀	中断后从最后一个完整日继续，无需重跑整个起报
NetCDF deflate-1 压缩	将全年输出从约 1.21 TB 降到约 600—700 GB
.nc.part 后原子替换	把中断时的半截文件排除在完成结果之外
用 UTC 日期生成标签	避免容器时区导致预报日期整体偏移一天
在 CUDA 上下文建立前创建进程池	避免子进程继承 CUDA 上下文造成挂起
ONNX 输出显式转为 float32 再做残差相加	避免 FP16 原地加法的类型转换错误
修正 long_name 属性	避免官方脚本全局替换留下的 longitudeg_name 拼写

断点恢复时，脚本读取最后一个完整 NetCDF 输出，重新归一化为内部状态，再继续后续预报日。由于写出的状态是反归一化后的 float32，恢复会引入极小的舍入误差，此前已完成的 lead day 无需重新计算。

复现中最容易踩的坑

1. ERA5 强迫日期错开一天

初始场为日期 $D$ 时，第一步读取 forcing.isel(time=0)，把 $D$ 日状态推到 $D+1$。因此 10 日强迫的正确区间是 $[D,D+9]$。

2. 深度轴方向正确，索引才有意义

23 层索引建立在“表层在前”的前提上。必须先排序深度，再按索引选择，并对照真实深度值；只看最终有 23 层并不能发现错误。

3. 不要提前改变 ERA5 单位

inference.py 内部会执行 ssr/3600、strd/3600 和 mtpr/1000。预处理侧应保留它期待的 ERA5 原始单位，否则会重复换算。

4. 经度插值必须考虑周期性

ERA5 和 GLORYS 都是规则经纬网格，但 $-{180}^{\circ }$ 与 ${180}^{\circ }$ 实际相邻。普通线性插值若没有周期补边，日界线附近会出现 NaN 或明显跳变。

5. GLORYS 与 GLO12v4 的掩膜不能混用

模型仓库提供的是 GLORYS 海陆掩膜。若把初始场换成业务化的 GLO12v4 分析预报产品，需要同步替换掩膜，否则海岸线差异会污染输入和输出。

6. CPU 可运行不等于可完成

ONNX Runtime 找不到 CUDAExecutionProvider 时会回退到 CPU，但全球 $1/{12}^{\circ }$、93 通道的自回归模型在 CPU 上通常不具备实际可用的吞吐量。日志中出现 CPU provider 时应先修环境，不要继续提交全年任务。

最终得到的复现链路

这次整理后，问海模型的复现成为一个分层、可恢复的工作流：

公开再分析资料
  -> GLORYS 23 层初始场 + ERA5 逐日强迫
  -> 形状、单位、深度、日界线和物理范围校验
  -> bulk formulae 计算 8 通道海气通量
  -> WenHai ONNX 残差式逐日预报
  -> 压缩 NetCDF、按日落盘、断点续跑
  -> SST / SSH 范围与后续统计验证

复现中最耗时的部分是准备约百 GB 的输入、管理数百 GB 的输出，并保证深度方向、时间约定、单位和掩膜都与训练时一致，加载模型本身只占很小比重。对于这类高分辨率地球系统 AI 模型，工程细节本身就是复现的一部分：维度校验通过只说明数组形状正确，物理量是否合理需要额外验证；任务能够启动只说明流程能跑起来，能否跑完并产出可信结果需要持续监控。