⚡

Earth2Studio诊断模型创建Skill earth2studio-create-diagnostic

该技能用于指导用户为NVIDIA Earth2Studio框架创建诊断模型包装器，涵盖简单派生诊断、AutoModel打包诊断与生成式/扩散诊断。包括模型类型分类、依赖配置、模型文件开发、坐标处理、前向实现、加载、测试、文档更新及验证流程。关键词：Earth2Studio、诊断模型、AutoModel、生成式诊断、气候预测、模型包装器、PyTorch、坐标处理、测试验证、NVIDIA。

气候预测 0 次安装 0 次浏览 更新于 9/6/2026
名称 earth2studio-create-diagnostic
版本 0.16.0
开源协议 Apache-2.0 metadata:
作者 NVIDIA Earth-2 团队 agent-skills@nvidia.com tags: [earth2studio, diagnostic-model, python]
描述 > 为单步数据转换创建 Earth2Studio 诊断模型包装器，包括简单派生诊断、打包的 AutoModel 诊断和生成式/扩散诊断。请勿用于预报时间步进模型、数据源或安装。 argument-hint: 参考推理脚本的 URL 或本地路径（可选）

快速入门检查清单

按顺序完成以下步骤，不要跳过。在编辑前完整阅读本 SKILL.md，并根据模型类型加载相关参考文件。实现后，先执行聚焦的 pytest 命令，再声称测试通过。若测试无法运行，请报告准确命令和失败原因，而不要声称成功。

  • [ ] 首先完整阅读本 SKILL.md
  • [ ] 获取参考脚本、仓库、论文或模型文档（步骤 0）
  • [ ] 将诊断模型分类为简单、AutoModel 或生成式（步骤 1）
  • [ ] 在编辑依赖文件前提交依赖额外包方案（步骤 1）
  • [ ] 创建 earth2studio/models/dx/<name>.py，只包含诊断 API
  • [ ] 创建 test/models/dx/test_<name>.py 并编写 mock 测试
  • [ ] 运行：uv run pytest test/models/dx/test_<name>.py -m "not package" -v
  • [ ] 必要时更新模型 extra、安装文档、API 文档和 CHANGELOG
  • [ ] 运行：make format && make lint && make license

关键命令规则：始终使用 uv run 执行 Python 命令：

  • 使用 uv run pytest ... 和 uv run python ...
  • 仓库工作流不要直接使用 pytest 或 python

如果生成的模型是错的，不要反复尝试同一修复。请遵循自我改进，修复本技能或其参考文件，然后继续使用修正后的工作流。

目的

实现一个诊断模型包装器，将第三方或派生的 ML 转换连接到 Earth2Studio。诊断模型在单个时间点转换数据：输入字段进，输出字段出，不做预报积分。

前置条件

  • 通过 uv 安装 Earth2Studio 并包含 dev 依赖（uv sync --all-extras）
  • Python 3.10+ 环境
  • 参考推理脚本、仓库、论文或模型文档
  • 打包模型的检查点来源和许可证信息

限制

  • 仅处理单步转换
  • 不支持时间步进预报模型；请使用 earth2studio-create-prognostic
  • 真实包测试可能需要访问 NGC、HuggingFace、S3 或其他注册表
  • 生成式验证可能需要 GPU 和固定种子用于有意义的比较

诊断模型类型

类型 继承 依赖 extra 示例
简单派生诊断 torch.nn.Module 通常不需要 Identity、风速
打包 AutoModel 诊断 torch.nn.Module, AutoModelMixin 必须有（即使为空） PrecipitationAFNO
生成式诊断 torch.nn.Module, AutoModelMixin 必须有（即使为空） CorrDiff

工作区

上下文 位置
Harbor 评估 写入 /workspace/output/earth2studio/models/dx/...
Harbor + --copy-repo 完整检出位于 /workspace/repo
本地克隆 包含 pyproject.toml 的目录

切勿读取 evals/targets/；这些文件仅供评分参考。

参考文件

根据需要按需加载：

文件 内容 加载时机
references/skeleton-template.py 简单、AutoModel、生成式包装的完整骨架 步骤 3-6
references/method-templates.py 坐标、加载、前向和设备方法的片段 步骤 4-6
references/testing-guide.py Mock、包、异常、样本和种子测试模式 步骤 7
references/validation-guide.md 参考比较、绘图、PR 卫生和审查跟进 步骤 10-11
references/pr-body-template.md PR 正文模板 步骤 11
references/pr-comment-template.md 验证评论模板 步骤 11

说明

步骤 0 — 获取参考资料

如果 $ARGUMENTS 提供了 URL 或本地路径，请使用。否则询问：

请提供参考推理脚本、仓库、论文或模型文档。

捕获参考模型的输入变量、输出变量、张量形状、归一化、网格、检查点来源、依赖要求和许可证。

步骤 1 — 分析类型并提出依赖

在编辑文件前分类所需诊断：

如果模型… 使用…
从输入计算派生量且不加载权重 简单诊断
从 Package 或外部检查点加载权重 AutoModel 诊断
产生多个样本/扩散输出/VAE 样本/随机超分结果 生成式诊断

依赖策略：

  • 简单派生诊断通常不需要 pyproject.toml extra。
  • AutoModel 和生成式诊断必须有一个命名可选依赖 extra，即使列表为空。
  • 在 [project.optional-dependencies] 下按字母顺序添加该 extra，并将其包含在 all 聚合中。
  • 在 OptionalDependencyFailure("model-extra") 和 @check_optional_dependencies() 中使用该模型 extra 名称。

在编辑 pyproject.toml 前，提交依赖 extra 方案并请求用户批准：

model-name = ["package1>=version", "package2"]
# 或者，当打包诊断不需要额外运行时包时：
model-name = []

步骤 2 — 添加依赖

批准后编辑 pyproject.toml：

  • 按字母顺序添加 extra。
  • 更新 all 聚合。
  • 优先使用参考包文档中支持的最低版本。
  • 除非参考模型确实需要并获批准，不要添加宽泛的不固定 Git 依赖。

步骤 3 — 创建模型文件

文件：earth2studio/models/dx/<lowercase>.py

使用现有模型文件中的仓库标准 SPDX/许可证头部。

简单诊断的常见导入：

from collections import OrderedDict
import numpy as np
import torch
from earth2studio.models.batch import batch_coords, batch_func
from earth2studio.utils import handshake_coords, handshake_dim
from earth2studio.utils.type import CoordSystem

打包和生成式诊断的额外导入：

from earth2studio.models.auto import AutoModelMixin, Package
from earth2studio.models.dx.base import DiagnosticModel
from earth2studio.utils.imports import OptionalDependencyFailure, check_optional_dependencies
from loguru import logger

标准方法顺序：

  1. __init__
  2. input_coords
  3. output_coords（用 @batch_coords() 装饰）
  4. 有用的 __str__
  5. AutoModel/生成式诊断的 load_default_package
  6. AutoModel/生成式诊断的 load_model
  7. 当非 PyTorch 状态需要移动设备时使用 to
  8. 私有/支持方法
  9. 用 @torch.inference_mode() 和 @batch_func() 装饰的 __call__

除非包装器自然拥有多个紧密相关的变体，且小型基类可以减少重复，否则避免共享基类或大型抽象。

步骤 4 — 实现坐标

诊断输入坐标通常使用公开的 Earth2Studio 顺序：

  1. batch：np.empty(0)，且在 OrderedDict 中第一项
  2. variable：输入变量名，使用 Earth2Studio 词汇表
  3. lat：公共纬度约定北到南，通常 90 至 -90
  4. lon：公共经度约定 0 至 360，端点通常不含

诊断包装器不应暴露 lead_time。如果诊断需要有效时间元数据，请在 coords["time"] 中记录为逐样本元数据；除非已有 dx 模式要求，否则不要将其作为张量维度。

output_coords 必须用 handshake_dim 和 handshake_coords 验证输入，然后更新输出变量，如有需要可更新输出 lat/lon 分辨率。生成式诊断必须在 batch 之后添加 sample 维度。

步骤 5 — 实现前向过程

使用单步 __call__，绝不创建迭代器。在模型执行前验证坐标，然后返回 (output_tensor, output_coords)。

@torch.inference_mode()
@batch_func()
def __call__(self, x: torch.Tensor, coords: CoordSystem) -> tuple[torch.Tensor, CoordSystem]:
    output_coords = self.output_coords(coords)
    x = (x - self.center) / self.scale
    out = self.core_model(x)
    return out, output_coords

对于生成式诊断，遍历 batch 维度并对每个输入项生成 number_of_samples 个样本。如果参考实现支持种子采样，请使用显式种子以确保可复现性。

步骤 6 — 实现模型加载

对于打包诊断：

  • load_default_package 应将 HuggingFace URL 锁定到提交（hf://org/repo@commit）或将 NGC/S3 版本锁定到不可变发布。
  • load_model 应调用 package.resolve(...)，先在 CPU 上加载检查点，将模块设为 eval()，并在适当处禁用梯度。
  • 仅在需要加载 pickle 的完整 PyTorch 对象时使用 weights_only=False。
  • 用 @check_optional_dependencies() 装饰可选模型类和 load_model。
  • 在 earth2studio/ 内使用 loguru.logger 记录加载消息，不要用 print()。

步骤 7 — 编写测试

文件：test/models/dx/test_<name>.py

必需测试：

函数 目的
test_<model>_call 使用 mock 或简单模型进行前向传递
test_<model>_exceptions 非法的坐标顺序、值或变量应抛错
test_<model>_package 使用 @pytest.mark.package 的 AutoModel/生成式诊断真实权重测试

生成式诊断还需要样本计数和确定性种子测试。请使用 references/testing-guide.py。创建一个 Phoo<ModelName> 哑模型，它匹配真实核心模型接口并产生确定性输出。

运行聚焦测试：

uv run pytest test/models/dx/test_<name>.py -m "not package" -v
uv run pytest test/models/dx/test_<name>.py::test_<model>_package --package -v

不要省略打包模型的包测试。如果任意随机输入对真实检查点物理上无效，请构造一个稳定的、适合模型的输入，同时仍加载真实权重并执行前向传递。

步骤 8 — 注册模型

对于公共模型，按字母顺序更新 earth2studio/models/dx/__init__.py。仅当用户明确要求不应导出内部或实验文件时才跳过注册。

步骤 9 — 文档

对于公共模型：

  • 按字母顺序添加到 docs/modules/models_dx.rst，使 API 文档包含生成的页面
  • 如果存在模型 extra，添加到 docs/userguide/about/install.md。包含模型说明，以及 pip install earth2studio[model-name] 和 uv add earth2studio --extra model-name 两种安装方式。
  • 在 CHANGELOG.md 的 ### Added 下更新。

格式化与 lint：

make format && make lint && make license

步骤 10 — 验证（如请求）

按照 references/validation-guide.md。创建未提交的原始、Earth2Studio、比较和健全性检查脚本。不要提交生成的输出、检查点、图像或本地验证脚本。

对于生成式诊断，固定种子并比较匹配的样本，或在无法完全相等时报告统计/容忍度一致。在继续前要求用户目视检查绘图。

步骤 11 — PR（如请求）

按照 references/validation-guide.md 并使用：

  • references/pr-body-template.md
  • references/pr-comment-template.md

在创建 PR 前，验证依赖 extra、all、安装文档、API 文档、变更日志、测试和验证产物是一致的。不要在 PR 文本中包含机器名、主机名、绝对路径、缓存路径、设备清单或上传的图像链接。使用绘图占位符供手动上传。

示例

简单诊断

用户：创建一个根据 u10m 和 v10m 计算风速的诊断模型。
助手：阅读 SKILL.md，分类为简单诊断，创建仅含 torch.nn.Module 的 windspeed.py，编写调用和异常测试，运行聚焦 pytest。

AutoModel 诊断

用户：从此参考脚本添加一个降水估计器。
助手：阅读 SKILL.md 和参考文件，提出依赖 extra，创建 torch.nn.Module + AutoModelMixin 包装器，编写 mock/包测试，更新文档/CHANGELOG/依赖，并运行验证命令。

生成式诊断

用户：包装一个扩散超分辨率模型。
助手：分类为生成式，添加样本输出坐标，支持种子处理，编写样本和确定性种子测试，并准备种子验证比较。

故障排除

错误 解决方案
OptionalDependencyFailure 使用 uv sync --extra <model-extra> 安装或修复 extra 名称
坐标握手失败 检查 OrderedDict 顺序和 handshake_dim 索引
输出形状错误 验证 output_coords 的长度是否匹配返回的张量形状
ModuleNotFoundError: pytest 使用 uv run pytest，不要用裸 pytest
包测试在随机输入上失败 使用稳定且物理上合理的输入，同时仍加载真实权重

提醒

应该：

  • 所有 Python 命令都使用 uv run python 和 uv run pytest。
  • 在 output_coords 上使用 @batch_coords()。
  • 在 __call__ 上使用 @torch.inference_mode() 和 @batch_func()。
  • input_coords 中始终把 batch 放在第一位，并设 np.empty(0)。
  • 用 handshake_dim() 和 handshake_coords() 验证坐标。
  • 在生成式 output_coords 中添加 sample。
  • 在每个 Python 文件带上仓库标准 SPDX/许可证头部。
  • 在 earth2studio/ 内使用 loguru.logger，绝不要用 print()。

不应该：

  • 继承 PrognosticMixin。
  • 包含 lead_time 坐标。
  • 创建 create_iterator。
  • 在没有明确多变体需求时，仅为一个包装器创建通用基类。
  • 提交 API 密钥、凭据、验证脚本、绘图或生成的输出。
  • 读取 evals/targets/。

自我改进

如果该技能产生错误输出，请在继续前更新它：

  1. 在生成代码或工作流中确定问题。
  2. 编辑 SKILL.md 或 references/ 中的相关文件以纠正指导。
  3. 为变更的技能文件运行聚焦验证。
  4. 若当前分支预期提交，请单独提交技能修复。
  5. 使用修正后的工作流继续模型实现。