Kornia 图捕获兼容性调查：docs/export_support 背后的 ONNX、torch.compile 与 torch.export 支持页机制

• 计算机视觉
• 人工智能
• 深度学习
• 图像处理

Kornia 文档站中的 "ONNX, torch.compile and torch.export support"（导出支持）页面，并非人工维护的表格，而是由docs/export_support/目录下的一套可复现的"全库图捕获调查"流水线自动生成的。每个 Kornia 公开算子都会被以具体 CPU float32 输入跑过 PyTorch 的三条图捕获路径（ONNX 导出、torch.export、torch.compile），逐条记录成功/失败及根因，最后合并成一份提交到仓库的 JSON 快照，在文档构建时渲染为网页。本文将讲解：测量了什么、case(...)注册协议、三条探测管线各自怎么工作、如何重新生成快照、合并阶段的校验策略，以及如何新增一个 case 并让它在页面上呈现。读者读完即可复现这份调查，或把同样的"注册-探测-合并-渲染"模式迁移到自己的 PyTorch 项目上。

背景链接：export_support 文档、导出支持页面（源）、页面入口（第 74 行引用该页面）。

一、整体架构：一份提交的快照，一次构建期渲染

docs/source/get-started/export-support.rst是最终支持页面的源文件，但它由 docs/generate_export_support.py 在文档构建时从 docs/source/_data/export_support.json 渲染生成。该 JSON 是docs/export_support/目录测量结果的提交快照；文档构建期间不会执行任何探测代码（"nothing here runs during a docs build"）。

流水线概览：

docs/export_support/cases_*.py # case 注册（每个公开算子 + 具体输入） │ run.py（并行调度 + 可续跑） ├── harness.py ──► results/onnx_*.json （ONNX 探测） └── probe_compile.py ──► results/compile_*.json （export + compile 探测） │ merge.py（校验 + 合并） ▼ docs/source/_data/export_support.json # 提交的快照 │ docs/generate_export_support.py（构建期渲染） ▼ docs/source/get-started/export-support.rst

harness.py 的模块 docstring 明确说：一个case是一个公开 kornia 可调用对象（函数或nn.Module实例）加上一组具体的张量输入；每个张量输入都会成为 ONNX 图的一个输入，Python 侧的关键字参数则烘焙为常量写入图中。

二、测量什么：6 个 case 注册文件

每个case是一个公开可调用对象 + 具体 CPU float32 输入，注册在 6 个cases_*.py文件之一：

快照头部（export_support.json末尾）记录了精确的测量环境，例如当前快照：

generated_at: 2026-09-02, kornia: 0.9.0rc1, torch: 2.9.1+cu128, onnx: 1.21.0, onnxruntime: 1.23.2, onnxscript: 0.5.7, opset: 18, python: 3.11.14, device: cpu, revision: e0969438

run.py 中GROUPS = ["aug", "feature", "geomA", "geomB", "misc", "models"]与上述文件一一对应。

从 cases_geomB.py 可以看到实测案例的典型形态，例如：

_cc("rad2deg", C.rad2deg, [torch.rand(2, 3) * 6.0]), # geometry.conversions _ce("find_fundamental", E.find_fundamental, [pts1, pts2], note="8POINT"), # geometry.epipolar _cam("PinholeCamera.project", lambda intr, extr, p: PinholeCamera(...).project(p), [...]) _clg("So3.exp.matrix", lambda v: So3.exp(v).matrix(), [so3_v])

同一个算子有多个代码路径时，用name[variant]约定命名（如find_fundamental[7POINT]、unproject_points[normalize]、RANSAC[fundamental_7pt]），使它们在页面同一行下分组；[random]后缀表示"采样在图中进行"的随机模式（详见第四节）。

三、case 注册协议：case(...)的参数

case(...)定义在 harness.py，签名如下：

def case( name: str, # 唯一显示名（通常用公开可调用名，如 "filters.filter3d"） group: str, # 报告分组（"filters"、"geometry.epipolar"、"feature.descriptors"...） target: Any, # nn.Module 实例或普通可调用 fn(*inputs, **kwargs) inputs: list[torch.Tensor], # 张量；每个都会成为 ONNX 图的活输入 kwargs: dict | None = None, # Python 侧关键字参数，烘焙为常量 *, note: str = "", # 页面上展示的自由文本 check: bool = True, # False：图内有随机性，只做导出+运行，输出不校验 atol: float = 2e-4, # 数值比较容差 rtol: float = 1e-3, skip: str | None = None, # 给出理由时该 case 直接标记 skipped，不运行 tags: tuple[str, ...] = (), # 例如 ("3d", "model", "pretrained") method: str | None = None, # 在模块上调用该方法而非 forward（如 "detect"） ) -> dict[str, Any]

关键语义（来自 docstring 与run_case实现）：

• inputs中的每个张量都是 ONNX 图的输入；非张量位置参数必须放进kwargs（前提是可调用对象支持按关键字接收），否则用 lambda 包裹；
• check=False用于图内带随机性的算子——RNG 流在图内不同，无法逐值对比，探测结果记为ok-unverified（程序能跑、输出形状正确且值有限）；
• skip用于不适合作为张量图的条目，例如cases_geomB.py中的save_pointcloud_ply（文件 I/O）、quaternion_to_angle_axis（已废弃别名），它们会以skipped/n/a呈现；
• method用于调用模块的非 forward 方法，如RANSAC.estimate_model_from_minsample（见cases_geomB.py中method="estimate_model_from_minsample"的用法）。

注册文件的主入口统一为（如cases_aug.py末尾）：

if __name__ == "__main__": run_cases(CASES, sys.argv[1], only=sys.argv[2:] or None)

即：python cases_aug.py <out.json> [names...]，可选地只跑部分 case（按名字、组或前缀过滤，见run_cases的only逻辑）。

四、三条探测管线

对每个 case 依次执行 eager 参考运行 + 三条捕获路径。

4.1 ONNX 探测（harness.py 的 run_case）

harness.py 的流程：

1. eager 参考：make_wrapper把任意可调用包装成固定元数的nn.Module（动态生成forward(in0, ..., inN-1)），并用_flatten_outputs把 kornia 的容器类输出（Boxes、Keypoints、Se3、Quaternion、Image等通过.data属性暴露张量）展平成张量元组；然后torch.manual_seed(SEED)（SEED=1234）下no_grad运行得到参考输出；
2. 导出：torch.onnx.export(m, inputs, dynamo=True, opset_version=18)，输出重定向到os.devnull以避免 dynamo 打印整张 FX 图；支持KORNIA_SURVEY_OPTIMIZE=0环境变量关闭导出优化；
3. 校验：onnx.checker.check_model(model, full_check=True)；
4. 运行对比：onnxruntime.InferenceSession(blob, providers=["CPUExecutionProvider"])运行图，与 eager 输出做np.allclose（close函数还处理 NaN 模式、bool 输出、形状差异），默认atol=2e-4, rtol=1e-3；
5. 记录：算子类型集合（op_types，含子图递归与自定义 domain 标记）、模型大小size_kb、墙钟时间time_s，失败时记录根因错误（err_str沿__cause__/__context__链取最深错误）以及涉事的 kornia 源码帧（_kornia_frames，最多 6 帧）。

状态全集（docstring 中列出）：ok | ok-unverified | eager-fail | export-fail | checker-fail | ort-load-fail | ort-run-fail | mismatch | no-tensor-output | skipped。

随机算子的处理：if not c["check"]时直接记ok-unverified，只验证输出个数与形状，不逐值比较——因为图内 RNG 流与 eager 不同。cases_aug.py中Rand包装器还重写了eval()返回 self，确保 harness 调.eval()后增强模块仍保持 train 模式（采样留在图内）。

4.2 torch.export 与 torch.compile 探测（probe_compile.py）

probe_compile.py 复用同一批注册文件（python probe_compile.py <cases_module> <out.json> [names...]），每条记录：

• export：torch.export.export(wrap, inputs)（非严格模式默认）→ok | fail | eager-fail | skipped；带 420 秒SIGALRM超时保护；
• compile：先用torch._dynamo.explain(wrap)(*inputs)统计graph break 数与 break 原因，再torch.compile(wrap)（默认 inductor 后端，fullgraph=False）运行并与 eager 对比 →ok | ok-breaks:N | ok-unverified | mismatch | fail | eager-fail | skipped；
• 额外记录目标可调用对象的可导入限定名（qualname），供合并阶段生成文档交叉引用；
• torch._dynamo.config.cache_size_limit = 64、suppress_errors = False。

两类特殊的进程级故障处理值得一提：

• Timeout：POSIX 上用signal.alarm，超时即Timeout异常；
• Poisoned：某个 case 把 dispatcher 搞坏（PythonDispatcherTLS未设置）后，解释器不可再用，进程以退出码POISON_EXIT = 3退出，由 run.py 负责重启续跑。

4.3 失败根因的归因

merge.py 把底层错误消息映射为页面上用户可读的"原因"，几个典型类目（_CAUSES表）：

• 数据依赖："Could not extract specialized integer from>pixi run -e default install-docs && pixi run -e default build-docs # 先构建一次文档，让算子获得交叉引用 python docs/export_support/run.py # ~2 小时（8 CPU 核）；可续跑 pixi run -e default build-docs

  run.py 的调度语义

  run.py 为每个 case 组启动两个并行子进程（ThreadPoolExecutor，共len(groups)*2个 worker）：onnx:<group>跑cases_<group>.py，compile:<group>跑probe_compile.py。支持的关键参数：

  参数	作用
  python docs/export_support/run.py	全量探测全部 6 组（约 2 小时）
  run.py aug misc	只跑指定 case 组
  run.py --merge-only	不做探测，仅用既有 results 重建 JSON
  run.py --force	即使结果不完整也强制合并

  可续跑：每个探测进程在results/下写onnx_<group>.json、compile_<group>.json及对应inventory_*.json；每条记录落盘后，重启时跳过已有记录的 case，因此崩溃（segfault 会留下crashed标记）的 worker 直接重启即可续跑（run_cases中先写 crash 标记再跑，run.py中MAX_RESTARTS=50循环重启被"毒化"的解释器）。

  版本印章（provenance）：inventory_*.json记录该探测用的 kornia revision（git rev-parse --short HEAD+-dirty后缀）与torch.__version__（harness.provenance）。prepare_resume发现记录来自不同 revision/torch 时丢弃而非续跑。

  merge.py 的合并校验

  merge.py 在写出快照前做三重校验（_validate），防止部分运行静默低估覆盖率：

  1. 每个 results 文件必须存在对应inventory_*.json；
  2. inventory 点名的每个 case 都须有非 crash 标记的记录；
  3. 所有组的 inventory 必须指向同一 kornia revision 与 torch 版本。

  任一不满足即中止合并并打印原因；--force可强制合并并在头部盖章"测量于什么版本"。merge.py还从最近一次文档构建的 Sphinx inventory（docs/build/html/objects.inv）解析交叉引用，使页面上的算子名成为可点击的文档链接；无 inventory 时算子以纯文本渲染。

  合并输出export_support.json每条 case 的字段包括：package、section、operator、variant、ref（最短文档化别名）、note、三列的*_detail原因文本、graph_breaks、where（首个 kornia 源码帧位置）。

  六、案例实操：以 augmentation 为例理解两种模式

  cases_aug.py 展示了最复杂的注册形态，其 docstring 明确两部分：

  • PART 1（确定性，组augmentation*）：aug(img, params=params)，把采样出的参数张量作为 ONNX 图输入喂入（Det包装器）；容器（AugmentationSequential等）与自动策略（RandAugment等）先 eager 调用一次，导出时把seq._params固定为常量；
  • PART 2（随机，组*.random）：aug(img)且保持 train 模式，随机性留在图内，check=False。

  以build()辅助函数为例，它为一个变换同时注册"确定性"与"随机"两个 case：

  build("Normalize", G2, lambda: KA.Normalize(mean=..., std=...), [IMG], direct=True) # direct=True：无随机参数 -> 确定性 case 就是普通 aug(*data) 调用 build("RandomRotation90", G2, lambda: KA.RandomRotation90(times=(1, 3), p=1.0), [IMG]) # 默认注册两个 case：RandomRotation90（params 喂入）+ RandomRotation90[random]（图内采样）

  值得注意的细节：

  • torch.manual_seed(7)在forward_parameters前固定种子，保证参数可复现；
  • 非张量参数（batch_prob、forward_input_shape）从params中剔除（NONPARAM_KEYS），保持为常量；
  • build(f"{_n}[p=0.5]", ...)用p=0.5的默认概率注册 5 个代表性子集（P05），使"逐样本 batch_prob 门控"分支真正进入图（这是常见的导出阻塞点，见P05注释）；
  • 容器用seq_build：先 eager 调用填充seq._params，导出时SeqDet.forward以params=params常量调用；
  • 页面上的变体分组（[variant]）由merge._split_name解析：RandomCrop[p=0.5][random]会变成 operator=RandomCrop、variant=p=0.5, random。

  七、页面如何渲染：状态语义与统计口径

  docs/generate_export_support.py 构建期渲染：

  • 状态 → 单元格（STATUS表 +_cell）：ok=“yes”；ok-breaks=“yes, N breaks”（可编译且匹配 eager，但有 graph break）；ok-unverified=“yes (random)”（随机算子，只验形状与有限性）；mismatch；fail=“no”（Details 列说明原因）；n/a=不适用的非张量图（文件 I/O、枚举、Python 容器、stub）；
  • 统计口径：每列计数为supported / probed（排除n/a行）；"distinct operators"按(package, operator)去重，且任一配置成功即计为 ONNX 支持（_operator_counts注释明确：[random]失败但确定性变体成功的增强算子算作支持，因为想导出 ONNX 的用户会用确定性变体）；
  • 页面布局：顶部三张统计卡片 + 图例（Legend）+ 按包汇总表 + 每包分节的算子表，附带一个隐藏的 JS 搜索/筛选控件（kornia-compat-controls：可按任意失败、ONNX 失败、graph breaks 等筛选行）；
  • 页脚有一段"失败是特定 torch/exporter/kornia 版本下的事实，而非承诺"的说明，并链接到 onnx 与 gpu-acceleration 文档。

  八、如何新增一个 case

  按 README 的 "Adding a case" 一节，在匹配的cases_*.py中用case(...)注册（协议见第三节，完整 docstring 在 harness.py）：

  • 唯一名称、用作页面分节的组名、可调用对象、张量输入、Python 侧 kwargs、页面上展示的 note；
  • 同一算子的变体使用name[variant]约定（如warp_perspective[fill]），合并时自动归到同一行标签下（_split_name会把组前缀剥离、把[a][b]折叠成 variant）；
  • 注册后重新跑python docs/export_support/run.py <group>（只跑该组，快，且only过滤可进一步限制到单个 case），再--merge-only合并，最后重建文档页面。

  新增 case 时若目标算子含随机性，记得check=False并给note说明；若无法以张量图表达（文件 I/O、枚举、stub），用skip="原因"让页面以n/a呈现而不是误报为失败。

  九、小结：这套机制的可迁移价值

  docs/export_support/的价值不止于一张支持矩阵，更在于它是一套可复现、可续跑、防静默失真的兼容性回归流水线：

  1. 注册式覆盖：新增算子只需一行case(...)，自动进入三条捕获路径的探测；
  2. 防失真：版本印章（revision/torch）+ inventory 校验 + 合并前完整性检查，杜绝部分运行覆盖完整快照；
  3. 可诊断：每个失败都带根因（数据依赖、缺 lowering、torch 版本 bug）与涉事 kornia 源码帧；
  4. 构建期渲染：快照提交进仓库，页面在文档构建时生成，数据跟着库走而不是手工维护。

  对任何希望系统化回答"我的 PyTorch 库在 ONNX / torch.export / torch.compile 下到底哪些能用"的团队，这套"cases 注册 → 并行探测 → 校验合并 → 构建期渲染"的模式都可以直接借鉴；对本仓库而言，它让用户在使用 ONNX 导出、部署或torch.compile加速前，先对某个算子的兼容性有一个版本化的事实依据。

  赞分享
  • 计算机视觉
  • 人工智能
  • 深度学习
  • 图像处理

  【免费下载链接】kornia

  🐍 Geometric Computer Vision Library for Spatial AI

  项目地址：https://gitcode.com/gh_mirrors/ko/kornia

  点击查看免费下载

  相关推荐

  • Kornia ONNX / torch.export / torch.compile 支持页面：从图捕获调研到自动生成的兼容性文档

    Kornia ONNX / torch.export / torch.compile 支持页面：从图捕获调研到自动生成的兼容性文档 本指南围绕 Kornia 仓

    计算机视觉人工智能深度学习图像处理

  • Kornia 全库算子 ONNX / torch.compile / torch.export 支持矩阵：图捕获调查与支持页面生成全解析

    Kornia 全库算子 ONNX / torch.compile / torch.export 支持矩阵：图捕获调查与支持页面生成全解析 Kornia 仓库自

    计算机视觉深度学习人工智能图像处理

  • kornia 全面支持 Dynamo ONNX 导出：约 140 个新增算子案例的图捕获兼容实践

    kornia 全面支持 Dynamo ONNX 导出：约 140 个新增算子案例的图捕获兼容实践 本文基于 kornia 仓库 changelog.d/+mig

    计算机视觉人工智能深度学习图像处理
  上一篇：FATE联邦学习部署终极教程：单机与集群完整方案指南 🚀

  下一篇：Apptainer 项目常见问题解决方案