WeatherNext实战：从环境配置到批量推理的完整指南

WeatherNext 是 Google DeepMind 在天气预测方向上公开的仓库名称。我最早关注它倒不是被“AI 预测天气”这个概念吸引而是想验证一个很现实的问题普通开发者如果只依赖个人电脑和公开气象数据到底能不能把这类机器学习天气预报项目跑起来跑通之后又怎么判断结果靠不靠谱。这个流程涉及的知识点不只是 PyTorch 或 GPU还包括气象数据格式、变量名、格点维度、批量任务设计和预报评估指标缺一环都会卡住。文章按实际落地顺序拆WeatherNext 到底在预测什么、运行前要准备哪些环境与数据、如何从单条样例跑到批量推理、用什么指标判断输出质量以及真实环境里最容易踩的坑。适合三类人看刚接触 AI for Science 方向的学生、想把机器学习预报模型接入业务系统的算法工程师以及需要评估预报工具的行业人员。1. 先搞懂 WeatherNext 预测的是什么别把它当天气 App1.1 输出是格点气象场不是一句“明天下雨”很多人第一次接触这类项目会下意识以为它像一个天气 App输入城市名返回“明天有雨”或“最高气温 28 度”。WeatherNext 这类模型并不是这么工作的。它处理的是一大片空间范围内的格点气象数据输出通常是一个多维数组维度包括起报时间、预测时效、经纬度以及多层气象变量。举个例子如果你向模型提供一个历史时间段的全球或区域气象场模型会给出未来 24 小时、48 小时甚至更长时间窗口内的预测结果。这些结果不是城市级别的文字描述而是覆盖整个网格的数值场比如 850 百帕温度场、海平面气压场、近地面风场。要得到具体站点的天气信息还得做空间插值和时间对齐。新手第一次看到输出文件时最需要建立的概念就是预测结果是“一张图”不是一个数字。判断模型好坏也不能只看某一个点的温度准不准而是看格点场整体的空间结构和时间演变是否合理。1.2 学习方法与传统数值天气预报的差异传统数值天气预报把大气运动方程离散到网格上用超级计算机不断迭代积分。它的优点是物理过程清晰可解释性强缺点是计算量非常大对算力、数据同化系统要求很高。WeatherNext 这类机器学习教育预测模型换了一条路从大量历史再分析资料中学习大气的时空演变规律然后用训练好的模型权重直接从输入气象场映射到未来气象场。这个差异带来的直接影响是推理速度。传统模式跑一次全球 10 天预报往往需要在超算集群上运行不短时间而机器学习模型在 GPU 上跑一次前向推理通常只需要几秒到几分钟。但速度快的代价是依赖数据分布。如果某个区域、季节或极端天气形态在训练数据里出现得很少模型的外推能力就会下降。所以不要把“速度更快”直接理解为“效果更好”。在熟悉的历史天气类型上ML 模型可能表现很好但面对极端事件或气候突变时仍需要额外验证。1.3 适合用机器学习预报模型的场景从落地角度我认为 WeatherNext 这类项目适合几个场景快速生成多成员集合预报用不同初始扰动得到一组结果用来估计预报的不确定性。对历史个例做批量回算评估模型在不同年份、不同天气系统下的表现。在算力资源有限的服务器上做短期预报实验作为传统数值预报结果的辅助参考。不适合的场景也很明确正式航空、航海、应急决策等对流程规范、可解释性和业务标准有严格要求的场景不能直接把模型输出当作最终决策依据。除非你已经完成了本地化验证并建立了完整的质量评估流程。2. 运行前先对待数据集和依赖版本比调参更关键2.1 硬件条件不要只看显存很多人以为跑深度学习天气模型必须先有一张超大显存的显卡。实际上模型推理阶段的显存要求不一定很高真正吃资源的是输入数据和数据预处理。我建议的最低配置可以是这样的思路显存至少 8GB内存至少 16GB磁盘预留至少几十 GB具体取决于你要处理的数据范围。如果你的输入数据只有小区域、短时段磁盘需求会小很多但如果想跑全球范围、长时间跨度、多层气压变量的数据几十 GB 可能只是起步。还有一种常见情况电脑可能能加载模型但在读取 NetCDF 或 GRIB 文件时内存直接打满。因为气象数据文件是经过压缩的读取时会把变量解压到内存里。多个变量、多个时间切片叠在一起内存占用会快速增长。2.2 依赖库的常见组合WeatherNext 这类仓库的实现通常基于 PyTorch 或 JAX数据层多用 xarray、numpy、netCDF4、cfgrib 这些库。第一次搭建环境时不要直接在全局环境里装包建议创建一个独立环境。conda create -n weathernext python3.10 conda activate weathernext pip install xarray numpy netCDF4 cfgrib torch上面只是通用的基础示例不代表仓库固定依赖。实际使用时以仓库里的 requirements 或环境配置文件为准。不要忽略版本问题cfgrib和netCDF4这类库经常因为底层 C 库版本不一致而报错修复起来比模型代码更麻烦。2.3 输入数据字段名必须对齐另一个非常容易被忽略的问题是字段名。模型内部可能把温度变量命名为2m_temperature但你下载的数据里可能叫t2m或者反过来。模型不会自动理解这两个名字是同一个东西它只会根据变量名去索引数组。我一般会先把样例数据文件的变量列表打印出来再和仓库文档或示例配置里的变量表逐个对照。字段对齐之后再处理维度顺序。比如有些数据按time, level, latitude, longitude排列有些则可能是time, latitude, longitude, level。如果不统一模型前向推理时会得到完全错误的结果。import xarray as xr ds xr.open_dataset(sample.nc) print(ds)这一步能解决的问题比你改十次模型参数都多。3. 从一条样例推理开始先别急着训练和批量3.1 用最小输入跑通第一轮第一次使用仓库时千万不要直接下载几十年的数据也不要一上来就准备大批量预测。建议先用小区域、短时间片、少量变量跑一条样例。重点观察三个环节是否正常数据加载、维度变换、输出形状。如果输入包含 24 小时的历史场那就先把这 24 小时的数据提取出来做成一个单批次的输入。运行成功之后再去调整时间范围和空间范围。这样能快速定位问题是出在数据读取、字段整理还是模型前向计算。第一次跑通后我会花十分钟记录三个关键信息输入文件大小、单次推理耗时、输出文件大小。这三项数据直接决定后续批量任务怎么设计。单条推理如果耗时 10 秒那 100 条就是 1000 秒如果中间还有数据读取和文件写入总时间会更长。3.2 一个通用伪代码示例不同版本的仓库接口差异可能很大所以这里不写死具体函数只给一个可以迁移到大多数 ML 气象项目的处理思路import xarray as xr import numpy as np # 1. 读取输入文件 ds xr.open_dataset(sample.nc) # 2. 选择目标时间范围例如过去 24 小时 input_ds ds.sel(timeslice(2020-06-01, 2020-06-02)) # 3. 按模型要求的变量和层级排序 model_input input_ds[model_vars].transpose(time, level, latitude, longitude) model_input model_input.values[np.newaxis, ...] # 增加 batch 维度 # 4. 执行模型前向推理 # output model(model_input) # 5. 把输出还原为 DataArray方便后续计算和可视化这个流程里最容易出错的是transpose这一步。气象数据的时间和维度顺序非常灵活如果不做显式排序模型拿到的张量很可能和训练时的输入分布不一致最终输出会是乱序的。3.3 第一次跑完检查什么模型推理成功后不要只看到“没有报错”就认为任务完成。我一般会先检查输出数组的 shape 是否和预期一致比如时间步长、空间网格数、变量通道数。接着打开输出文件看一下是否存在大量 NaN。出现 NaN 的原因很多有时是输入数据有缺测有时是经纬度顺序错误导致计算越界有时是模型权重加载不完整。如果第一条就报错先不要急着改模型代码。先回退一步确认报错来自数据加载、维度转换还是模型前向。绝大多数新手问题出现在前三步。4. 批量推理时要设计队列、命名和重试机制4.1 不要一上来就多进程单条样例能跑通和批量任务能稳定运行是两个完全不同的概念。批量任务最容易出问题的点不是模型本身而是输入路径错误、内存缓慢增长、输出文件名冲突。我习惯先用一个 5 条输入的短列表测试确认输出目录下有 5 个文件而且每个文件的文件名都能对应到正确的时间戳。如果这个测试通过再扩大到 20 条、100 条。直接开多进程跑全量数据是很危险的做法因为一旦中途崩溃你很难知道哪些任务已经写入哪些还需要重跑。4.2 按起报时间切片组织输入批量任务建议按init_time也就是起报时间把数据切成独立的子任务。每个起报时间对应一条或多条预测任务输出文件名用起报时间做标识例如forecast_20200601_0000.nc。这种组织方式的好处是容错性强。如果一个时间点失败只需重新跑那一个任务不需要重算整个集合。后续做时间序列分析的时候也能很清晰地按起报时间对齐结果。4.3 加入重试和日志机制长时间在服务器上跑批量任务脚本必须记录每个任务的状态。我一般会用简单的 markdown 或 CSV 文件记录三列任务名、状态、报错信息。任务状态分为 pending、running、ok、failed 四类。遇到失败时可以写一个最多重试三次的循环。不要无限重试因为某些错误即使重试一百次也不会成功比如数据文件缺失或字段名错误。每次失败时需要把完整报错写入日志方便重建任务时直接定位原因。注意批量任务最怕的不是单条失败而是失败后继续覆盖写文件导致最终结果里混合了新旧数据。因此输出文件名建议包含输入时间戳并且使用独立输出目录避免多个任务写入同一个文件。5. 评估预报效果时不要只看“看起来像不像”5.1 常用定量指标WeatherNext 这类模型在论文和公开报告中常出现三个指标RMSE、ACC 和 CRPS。RMSE 用来衡量预测值和观测值之间的平均误差适合评估连续变量比如温度、位势高度。ACC 是异常相关系数用来判断预测的空间分布是否和实况一致数值越接近 1说明分布吻合度越高。CRPS 常用于集合预报它同时考虑预测精度和不确定性能反映出概率预报的综合能力。第一次评估时至少看 RMSE 和 ACC 两项。只看单一指标很容易得到片面结论。比如某个模型 RMSE 很低但空间分布整体偏平滑看起来误差小实际锋面位置完全不对。5.2 需要和一个简单基准对比单独看模型自己的 RMSE 没有意义。同一个时间段至少要和一个简单 baseline 对比比如 persistence 和 climatology。Persistence 的意思是未来与当前相同也就是把当前气象场直接当作预测结果Climatology 是用多年气候平均值当作预测结果。如果模型连这两个最简单的基准都不能稳定超过那说明模型很可能在数据拟合或输入处理上存在问题。只有当模型明显优于 persistence 和 climatology才有继续优化的价值。5.3 空间分布和极端事件单独评估平均指标会掩盖很多问题尤其是强降水和台风路径这类极端事件。建议按区域、季节、天气型分别统计指标。举个例子一个模型可能在温带气旋路径预测上表现很好但对热带气旋强度预测很差。如果只用一个总平均 RMSE 来汇报很难暴露这个问题。我处理类似项目时遇到过一种很常见的现象短期预报 RMSE 很好看但过了 48 小时后气压场出现明显平滑化极端低值被削弱。这是因为模型训练时倾向于生成“平均”的结果避免在一个不确定的时间点上给出极端预测。如果你不单独检查极端值区域很难发现这个隐患。所以判断模型能力不能只看它是否画出了一张“好看”的天气图而是看关键天气系统比如气压槽、锋面、台风中心位置以及极端值是否被保留。6. 实际运行中最容易踩的四个坑及排查顺序6.1 数据路径、字段名、坐标系如果模型输出全是 NaN或者预测结果明显空间错位先检查三样东西路径下的文件是否完整是否是空文件变量名是否和模型预期一致经纬度是否按模型要求的顺序排列。很多气象数据按照纬度从小到大排列也就是南到北。如果模型默认按北到南排列而你直接把数组喂进去输出结果会南北颠倒。这个问题从整体指标上很难一眼发现必须画图检查。6.2 依赖版本与 C 库问题netCDF4、cfgrib 这类库经常遇到 HDF5 或 C 库版本冲突。报错通常会提示找不到某个.so文件或者加载 libgrib 失败。解决办法不是修改模型代码而是重建干净环境。我的处理顺序通常是把当前环境删掉重新创建使用 conda 安装数据和科学计算依赖然后手工安装模型相关依赖。比起直接在原有环境里逐个升级包重建环境往往更快。6.3 显存大小和 batch size推理时出现 CUDA Out of Memory第一反应是降低 batch size。不需要急着换大显存的显卡。如果 batch size 已经降到 1仍然 OOM那就考虑裁剪输入区域或者在加载模型时启用 half 精度。另外注意一下输入张量的尺寸。如果输入数据是全球 1.0 度分辨率但模型训练时使用的是 0.25 度或特定重采样后的网格输入维度和你本地文件不一致也会导致内存异常。6.4 排查顺序现象、输入、环境、参数遇到问题不要一上来就调整预测时效或变量组合。完整排查顺序应该是先看现象是直接报错还是输出为空还是结果明显异常。再看输入文件路径、字段名、缺失值、时间范围和经纬度顺序。再看环境依赖版本、显存、内存、磁盘。最后再看参数batch size、预测时效、并发放置。日志和报错堆栈是最重要的信息源。一定要确认报错发生在数据加载、模型前向还是后处理阶段。很多问题表面上像模型能力不够实际上只是前置数据和环境没有处理好。我自己的建议是这类项目第一周不要追求训练。先花时间把推理样例和批量流程跑稳定等你对数据维度、评估指标和输出语义足够敏感之后再考虑训练和大规模验证。踩过几次坑之后你会发现真正卡住你的往往不是模型本身而是那些看起来很简单却容易被忽略的细节。