Earth2Studio 中文学习手册0.19.0交互实验 ↗
CHAPTER 04 / LEARN BY UNDERSTANDING

最小可运行例子

这一章提供两个层次:不下载模型的 CPU 冒烟测试,以及真实 FCN 推理。代码按 0.19.0 源码接口核对;这里没有声称已在某张 GPU 上完成实测。第一次运行应在独立目录,不能覆盖生产数据。

先准备独立环境

0.19.0 的 Python 范围是 >=3.11,<3.15。本教程选 Python 3.13;服务器优先 Linux。macOS 可以学习数据结构和运行部分 CPU 功能,但不能把 Apple GPU 当作 CUDA 环境。依赖定义

已有 uv 时:

mkdir earth2-lab
cd earth2-lab
uv venv --python 3.13
source .venv/bin/activate
uv pip install 'earth2studio==0.19.0'
python -c 'import earth2studio; print(earth2studio.__version__)'

基础包也包含科学计算依赖;“CPU 测试”表示不需要模型权重和 CUDA,并不表示零依赖。若 GRIB 依赖安装失败,检查 Python wheel、系统架构及 ecCodes;不要在业务服务环境里反复修改版本。官方安装指南

例子 A:不下载权重的流程测试

保存为 smoke.py。这里使用随机数组和保持初值不变的 Persistence,不是有气象意义的 AI 预报。

from collections import OrderedDict
from pathlib import Path
import numpy as np
import torch
from earth2studio import run
from earth2studio.data import Random
from earth2studio.io import ZarrBackend
from earth2studio.models.px import Persistence

domain = OrderedDict([
    ('lat', np.array([20.0, 10.0, 0.0])),
    ('lon', np.array([100.0, 110.0, 120.0, 130.0])),
])
path = Path('smoke.zarr')
if path.exists():
    raise SystemExit('请使用一个尚不存在的输出路径')
model = Persistence(['t2m'], domain)
io = ZarrBackend(str(path))
run.deterministic(
    ['2025-01-01T00:00:00'], 2, model,
    Random(domain), io, device=torch.device('cpu'),
)
values = io['t2m'][:]
assert values.shape == (1, 3, 3, 4)
np.testing.assert_array_equal(values[:, 0], values[:, 1])
np.testing.assert_array_equal(values[:, 1], values[:, 2])
print(io.root.tree())
print('通过:1 次起报,0/6/12 小时,3×4 网格')

运行 python smoke.py。这个例子验证模型接口、坐标、时间推进和 Zarr 写入。其 API 用法与上游确定性工作流测试一致。上游测试

例子 B:真实 FCN 的 24 小时预报

在有 NVIDIA GPU 的 Linux 环境中,先确保驱动和所装 PyTorch 的 CUDA 运行时相容;基础包安装成功并不保证 torch.cuda.is_available() 为真。按 PyTorch 安装选择器 选择对应构建,再安装:

uv pip install 'earth2studio[fcn]==0.19.0'
python -c 'import torch; print(torch.__version__, torch.version.cuda, torch.cuda.is_available())'
export EARTH2STUDIO_CACHE="$PWD/cache"

保存为 forecast.py:

from collections import OrderedDict
from pathlib import Path
import sys
import numpy as np
import torch
from earth2studio import run
from earth2studio.data import GFS
from earth2studio.io import ZarrBackend
from earth2studio.models.px import FCN

if len(sys.argv) != 3:
    raise SystemExit('用法: python forecast.py 起报UTC时间 输出目录.zarr')
if not torch.cuda.is_available():
    raise SystemExit('本例要求可用的 CUDA GPU')
init, destination = sys.argv[1:]
if Path(destination).exists():
    raise SystemExit('输出路径已存在,请换一个路径')
Path(destination).parent.mkdir(parents=True, exist_ok=True)
model = FCN.load_model(FCN.load_default_package())
print('输入坐标:', model.input_coords())
io = ZarrBackend(destination)
run.deterministic(
    [init], 4, model, GFS(), io,
    output_coords=OrderedDict({
        'variable': np.array(['t2m', 'u10m', 'v10m', 'msl'])
    }),
    device=torch.device('cuda:0'),
)
assert io['t2m'].shape[1] == 5
assert io['t2m'].shape[-2:] == (len(io['lat']), len(io['lon']))
print(io.root.tree())
print('时效:', io['lead_time'][:])
python forecast.py 2025-01-01T00:00:00 outputs/fcn-20250101.zarr

选择固定历史起报便于复现;该时刻位于 GFS 适配器的历史支持范围,但实际下载仍受上游服务及本机网络影响。若失败,应先辨别是数据不可达、缺变量还是 GPU 问题,不能自动换一个日期却保留旧标签。GFS 适配器

output_coords 只限制写出的变量,不会把模型内部变成只计算四个变量的小网络。模型仍需要完整输入状态;全局模型也不能通过只给香港小区域输入就自动省去全球计算。

如何确认结果看起来正常

先检查结构,再检查数值:起报唯一、时效 0/6/12/18/24 小时、坐标与数组大小一致;t2m 为温度量,u/v 为东西/南北风分量;没有非预期的 NaN 或 Inf。随后在若干已知经纬位置抽样,并绘制全球场观察日期线、极点与南北方向。

不要以“输出一个文件”作为成功标准,也不要把 min/max 落在某个宽泛范围当作预报准确的证据。

用容器缩短环境排查

主机需要 NVIDIA 驱动和 NVIDIA Container Toolkit。容器镜像包含用户态运行库,不能代替宿主机 GPU 驱动。Container Toolkit 安装说明

官方 0.19.0 安装说明列出 NVIDIA PyTorch 26.04-py3 作为一种环境起点。下面在此基础上显式请求 GPU,并将缓存与输出保留在宿主机:

docker run --rm -it --gpus all --ipc=host \
  -v "$PWD:/workspace/lab" -w /workspace/lab \
  -e EARTH2STUDIO_CACHE=/workspace/lab/cache \
  nvcr.io/nvidia/pytorch:26.04-py3

容器内先验证 torch 能访问 GPU,再按固定版本安装 FCN extra。安装系统依赖的具体命令见官方指南。此交互容器用于测试;生产应将成功环境构建成固定 digest 的镜像,避免每次启动联网安装依赖。

从 FCN 换到 FCN3

脚本中的类可以换为 FCN3,但环境不能只改一个 import。0.19.0 的 fcn3 extra 涉及 Makani 和 torch-harmonics,仓库 uv 配置还固定了部分依赖来源。为了保留这些约束,建议在新的环境中检出固定 tag,用仓库的锁文件安装所需 extra,再做 GPU 兼容性验证:

git clone --branch 0.19.0 --depth 1 https://github.com/NVIDIA/earth2studio.git
cd earth2studio
uv sync --frozen --extra fcn3 --python 3.13

这是一条保留仓库依赖来源的构建路线,仍要求满足系统编译和 CUDA 条件;不是任何 GPU 上一条命令必然成功的承诺。尤其注意 torch-harmonics 的 DISCO CUDA 扩展:包装代码在缺失时会提示 GPU 性能下降。是否使用 BF16、成员批大小和具体驱动组合应以该模型的验证结果确定。依赖来源、FCN3 加载实现