ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

GraphCast气象预报工具快速入门完全指南:4步在免费Colab上跑出全球预报

GraphCast气象预报工具快速入门完全指南:4步在免费Colab上跑出全球预报 GraphCast气象预报工具快速入门完全指南4步在免费Colab上跑出全球预报【免费下载链接】weathernext项目地址: https://gitcode.com/GitHub_Trending/gr/weathernextGraphCast 是 Google DeepMind 出品的深度学习气象预报工具它把图神经网络 GraphCast 与扩散模型 GenCast 封装成一套可直接运行的中期天气预报系统1~15 天全球尺度并且是唯一提供预训练权重、归一化统计量和示例输入数据的开源实现。本文带你先在免费 Colab 上 10 分钟内跑出第一份预报再对比 4 种部署环境、完成 3 个典型实战任务最后才回头拆解原理。一、快速上手最短路径跑通结果最轻量的路线是免费的 Colab gencast_mini_demo.ipynb它运行的是GenCast 1p0deg Mini 20191° 分辨率、13 个气压层、4 次细分二十面体网格也是四个官方模型中内存占用最小的一个能跑在 Colab 免费提供的 TPUv2-8 上你本地不需要任何硬件。步骤1克隆仓库并打开演示笔记本git clone https://gitcode.com/GitHub_Trending/gr/graphcast cd graphcast拿到仓库源码后将gencast_mini_demo.ipynb导入 Colab 打开后续所有操作都在浏览器里完成。步骤2把运行时切换为 TPU点击菜单栏运行时 → 更改运行时类型在硬件加速器一栏选择 TPUTPUv2-8免费。不选 TPU 的话只能用 CPU30 步的自回归外推rollout即把上一步输出反复喂回模型逐步推演会慢到无法接受。步骤3运行安装与初始化单元依次运行Upgrade packages、Pip install repo and dependencies两个单元会自动安装仓库代码与 JAX、Haiku、xarray 等依赖再运行Reconfigure jax if running on TPU单元刷新 TPU 运行时。运行完成后控制台不再报错即表示环境就绪。步骤4选模型、选数据并加载在Choose the model下拉框里选带Mini字样的 checkpoint预训练快照运行Load the model单元随后在Choose data to plot前一步的下拉框里选与模型匹配的 1° ERA5 示例数据集Notebook 已按分辨率自动过滤如source-era5_date-2019-03-29_res-1.0_levels-13_steps-30.nc运行Load weather data单元。加载成功后单元格下方会显示一个含time维度的 xarray 数据集NetCDF 格式的多变量大气场容器。步骤5构建模型并跑出自回归预报先运行Extract training and eval data与Load normalization data两个单元再运行Build jitted functions…最后运行Autoregressive rollout (loop in python)单元。跑完后你会得到predictions对象8 个集合成员 × 30 个 12 小时步约 15 天的全球预报场。✅ 验证点接下来运行Plot prediction samples and diffs单元能看到目标值 / 预测值 / 差值三栏动画随时间推进且温度场默认2m_temperature形态与目标基本吻合说明预报确实跑通了。Q: 第一个 rollout 单元跑了五六分钟没输出是卡死了吗 A: 不是。首次执行包含 JAX 的即时编译tracing compilation是一次性固定开销同一会话内再跑一遍会快得多。Notebook 中Run the model的说明段落也明确提示了这一点。二、环境选型一张表看清所有方案环境适用场景硬件要求安装/启动命令免费 ColabTPUv2-8学习体验、跑 GenCast Mini无本地要求完全免费Notebook 首个单元自动 pip 安装本地机器CPU读源码、改代码调试≥16GB 内存pip install -e .云 TPU v5e4 芯片1° 版 GenCast 推理约 21GB 主机内存 约 8GB HBMgcloud compute tpus queued-resources create node-1 --node-idnode-1 --zoneus-south1-a --accelerator-typev5litepod-4 --runtime-versionv2-tpuv5-litepod云 TPU v5p8 芯片0.25° 高分辨率业务级推理约 250GB 主机内存 约 32GB HBM同上改为--accelerator-typev5p-8 --zoneus-east5-a --runtime-versionv2-alpha-tpuv5云 GPUH100只有 GPU 资源的场景0.25° 版需约 300GB 系统内存 约 60GB 显存需把模型 attention 改为triblockdiag_mha见下方说明上图为官方对比同一份 0.25° GenCast 权重在 TPUv4 与 H100 上生成预报的误差差值RMSE 与 CRPS色标 ±3%。绝大多数要素误差差在 ±3% 以内H100 平均仅差约 0.3%RMSE和 0.4%CRPS精度基本无损但 0.25° 模型 30 步推理要 25 分钟左右比 TPU 的 8 分钟慢约 3 倍。云 TPU 的成本参考docs/cloud_vm_setup.md中给出v5e 4 芯片跑 1° GenCast 的 30 步 rollout 约 5 分钟含首次编译单次成本约 0.11~0.48 美元v5p 8 芯片跑 0.25° 版约 30 分钟成本约 2.50~11.10 美元。Spot抢占式实例还能再打 0.09~0.4 折。在 Console 里创建 TPU 时先确认项目如上图创建成功后 SSH 进 VM 启动 Jupyter 服务把生成的本地 URL 填进gencast_demo_cloud_vm.ipynb的连接对话框即可python3 -m notebook --port8081 --no-browser执行后终端会打印一个http://localhost:8081/...?token...的链接把它粘贴进 Notebook 的 Local Runtime 弹窗浏览器随即接管远端 TPU 上的内核。Q: 创建 TPU 时报Stockout错误怎么办 A: 说明该区当时没有空闲配额。勾选创建表单里的Enable queuing命令行方式则用queued-resources create请求会进入Compute Engine TPUs Queued Resources排队容量释放后自动变为 active。另外注意UI 里的pre-emptibility选项已废弃不要勾选否则会直接报错。三、实战演练完成 3 个真实任务3.1 用 GenCast 做 15 天全球集合预报数据准备Mini demo 自带示例数据source-era5_date-2019-03-29_res-1.0_levels-13_steps-30.nc含 2019-03-29 起 30 个 6 小时帧的 ERA5 再分析场。运行Extract training and eval data单元后eval_inputs取前 2 帧作为初始条件GenCast 使用当前帧 12 小时前的帧eval_targets是其后每个 12 小时时刻的真值eval_inputs, eval_targets, eval_forcings ( data_utils.extract_inputs_targets_forcings( example_batch, target_lead_timesslice(12h, f{(example_batch.dims[time]-2)*12}h), **dataclasses.asdict(task_config)))模型运行核心是rollout.chunked_prediction_generator_multiple_runs它按预报步迭代一步预测被 jit 编译、8 个集合成员 pmap 并行到多块 TPU 上num_ensemble_members 8 rngs np.stack([jax.random.fold_in(jax.random.PRNGKey(0), i) for i in range(num_ensemble_members)], axis0) chunks [] for chunk in rollout.chunked_prediction_generator_multiple_runs( predictor_fnrun_forward_pmap, rngsrngs, inputseval_inputs, targets_templateeval_targets * np.nan, forcingseval_forcings, num_steps_per_chunk1, num_samplesnum_ensemble_members, pmap_devicesjax.local_devices()): chunks.append(chunk) predictions xarray.combine_by_coords(chunks)跑完得到predictionssample维为 8、time维为 28 个 12 小时步约 14 天的集合预报。注意集合成员数应是 TPU 芯片数的倍数TPUv2-8 下用 8 正好一块芯片一个样本。结果可视化运行Plot prediction samples and diffs单元对任意要素默认 2 米温度逐样本输出Targets / Predictions / Diff三栏动画再运行Plot ensemble mean and CRPS单元得到集合平均场和逐格点的 CRPS连续排名概率得分衡量预报分布与真值偏差的综合指标热力图。踩坑说明集合成员数若不能整除芯片数pmap 会报设备分片错误首次运行包含编译时间别中途断开。3.2 用 GraphCast 做确定性逐步预报数据准备打开graphcast_demo.ipynb模型下拉框选 checkpoint官方提供 0.25° 正式版、1°GraphCast_small和 HRES 微调的 operational 版三个权重示例数据集覆盖 fake/era5/hres 三种来源、1°/0.25°/6° 等分辨率。GraphCast 以 6 小时为步长train_inputs, train_targets, train_forcings ( data_utils.extract_inputs_targets_forcings( example_batch, target_lead_timesslice(6h, f{eval_steps*6}h), **dataclasses.asdict(task_config)))模型运行GraphCast 是确定性模型无集合维rollout 用单跑接口predictions rollout.chunked_prediction( run_forward_jitted, rngjax.random.PRNGKey(0), inputseval_inputs, targets_templateeval_targets * np.nan, forcingseval_forcings)结果可视化Plot predictions单元同样给出目标 / 预测 / 差值三栏动画可以直观看到低压槽东移、对流云区演变的模拟过程。踩坑说明该 Notebook 在 rollout 前有一行分辨率断言——model_config.resolution必须等于360/lon格点数选错数据集会直接抛AssertionError按提示重选匹配模型的数据即可另外 0.25° 权重在纯 CPU 运行时内存不够请用 1°small权重或换 TPU。3.3 验证训练链路算损失与梯度数据准备沿用 3.1 的train_inputs/train_targets/train_forcings12 小时单步训练切片。模型运行loss, diagnostics loss_fn_jitted( rngjax.random.PRNGKey(0), inputstrain_inputs, targetstrain_targets, forcingstrain_forcings) print(Loss:, float(loss))再用grads_fn_jitted可得到全网络梯度Mean |grad|量级在 1e-3 附近为正常。结果可视化diagnostics是逐要素的损失分解 xarray可挑要素打印数值确认各通道都有有效梯度。踩坑说明损失/梯度比推理更吃内存——gencast_mini_demo.ipynb明确提示这些单元requires larger amounts of memory than running inference只在小模型Mini 或 random 架构上稳妥。Q: 运行 GraphCast 演示时 cartopy 绘图直接崩溃 A: Colab 内核自带的 shapely 版本与 cartopy 不兼容。graphcast_demo.ipynb中已内置Workaround for cartopy crashes单元先pip uninstall -y shapely再从源码重装 shapely跑一遍即可。核心模块路径/graphcast/rollout.py # 推理期自回归外推循环 /graphcast/gencast.py # GenCast 一步预测 采样器配置 /graphcast/graphcast.py # GraphCast 主架构四、原理深挖它是怎么工作的两个模型共享一个思路不要直接预报而是让网络学会大气状态如何演化一小步然后用循环把这一小步叠成 10 天。GraphCast 的类比是城市路网上的邻居互报。地球表面被铺上一层二十面体三角网格mesh比经纬网格在两极更均匀每个网格点是一个节点、特征向量装着温度/湿度/风场等几十项要素。模型先编码把网格数据搬到粗到细的多层网格图上再由一个 4 次消息传递的图网络做处理——每个节点反复向邻居要信息、更新自己的状态这恰好对应大气的局地耦合低压吸引周围气流这种物理过程靠消息传递步数来学习最后解码把网格上的状态投回规则经纬网格输出预报。整个编码器→图处理器→解码器在 graphcast/graphcast.py 里实现mesh_size、gnn_msg_steps、latent_size三个超参分别控制网格细化层数、消息传递轮数和节点特征宽度。GenCast 换了一台发动机它不输出最可能的那一个而是生成一个分布。类比修图软件里的去噪——从一团纯噪声出发用一个去噪器denoiser内部同样是网格上的稀疏 Transformer分 20 个噪声等级逐步擦掉噪声最终擦出一份 12 小时后的天气。每个 12 小时步生成 8 个样本就是集合预报样本之间的离散程度天然就是不确定性采样用 DPM-Solver 2S 加速器graphcast/dpm_solver_plus_plus_2s.py比朴素反向扩散快得多。另外它还会按历史统计量对输入做归一化、对输出按时间差分布做反归一化graphcast/normalization.py相当于让模型只学变化量而不用记绝对数值。五、性能与边界什么时候好用、什么时候别用上图为 Mini 模型8 成员集合对比 ECMWF 51 成员 ENS2019 年公平 CRPS的误差比蓝色比 ENS 好红色更差色标 ±20%。可以读出三条边界好用区500~850hPa 中低层、3~10 天的中期预报大面积为蓝尤其比湿q与 2 米温度在 5~10 天时效显著优于 ENS——这正是中期预报5~10 天这个用武之地。别指望的区200hPa 平流层附近的风场u/v在 0~3 天大面积偏红高层位势高度z前 7 天也偏弱超过 10 天后各要素误差普遍抬升。Mini ≠ 主力官方明确说 Mini 精度reasonable but not representative它是为低成本演示牺牲精度的版本业务级精度要看 0.25° 版论文中 0~10 天多数要素 RMSE 优于 ENS。优化方法适用场景效果复杂度0.25° 降到 1°/Mini内存受限≤21GB 主机内存即可跑 1°主机内存从约 250GB 降到约 21GB低换权重即可集合成员 50 → 8演示、学习、显存紧张内存与耗时约为 1/6低改num_ensemble_membersSpot 抢占式 TPU成本敏感的非关键任务折扣 60%~91%可能被随时回收低命令加--spotBFloat16 混合精度casting.pyTPU 上跑 GraphCast显存减半、计算更快低Notebook 已默认包裹GPU 上换triblockdiag_mhaattention无 TPU 只有 H100 类 GPU精度降约 0.3%~0.4%0.25° 推理约慢 3 倍中需改SparseTransformerConfig两个字段六、排错FAQ与资源导航Q: TPU 上跑 Notebook 时 JAX 行为怪异或导入libtpu失败 A: Colab 的 TPU 镜像里 jax/libtpu 版本偏旧。gencast_mini_demo.ipynb已内置处理单元pip uninstall -y libtpu libtpu-nightly后重装jax[tpu]运行它并重启内核。Q: 想用自己的初始场去哪里拿 ERA5 数据 A: 演示用的小样本已放在仓库配套的dm_graphcast公开 Cloud Bucket 中Notebook 直接匿名读取完整训练数据是 ECMWF 的 ERA5 再分析1979 年起、6 小时分辨率推荐按 Weatherbench2 的 Zarr 形式获取。注意数据时间格式必须与模型要求一致GenCast 用当前帧 12 小时前帧GraphCast 用 6 小时步。Q: 输出能用于商业产品吗 A: 分两份代码是 Apache 2.0可商用但预训练模型权重是 CC BY-NC-SA 4.0仅限非商业用途。做商业预报服务需要自行训练权重。Q: 本地 CPU 能跑吗 A: 能但很慢。graphcast_demo.ipynb支持选 random 小模型 6° 低分辨率数据在 CPU 上做 4 步 1° 推理勉强可行要完整体验请直接用 Colab 免费 TPU。资源导航云端部署文档docs/cloud_vm_setup.md含 TPU 选型、成本估算、GPU 备选方案演示笔记本gencast_mini_demo.ipynb免费 TPU、graphcast_demo.ipynbGraphCast 三权重、gencast_demo_cloud_vm.ipynb云 TPU源码入口graphcast/ 目录predictor_base.py定义统一接口losses.py含纬度加权损失论文GraphCast 发表于 Science 2023Lam et al.GenCast 为 arXiv:2312.15796Price et al.两者引文格式见 README.md从免费 Colab 的一次 rollout到 v5p 八芯片的业务级推理这条链路你已经完整走过一遍GraphCast 教你大气一小步怎么演化GenCast 在此基础上把预报升级为一个带不确定性的分布。建议的下一步实践先跑通 Mini demo然后把初始日期换成最近一年里的某一天对比 8 成员集合平均与你选要素的目标场 CRPS 曲线——这会让你对AI 预报的置信度建立起第一手直觉也为切换到 0.25° 高分辨率版本打好判断基础。【免费下载链接】weathernext项目地址: https://gitcode.com/GitHub_Trending/gr/weathernext创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表