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

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

【免费下载链接】kornia🐍 Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/GitHub_Trending/ko/kornia

本指南围绕 Kornia 仓库中新增的ONNX、torch.compile 与 torch.export 支持（Export Support）页面展开，介绍它是如何在构建时由一个提交到仓库的调研快照docs/source/_data/export_support.json自动渲染出来的，以及支撑它的整套图捕获（graph-capture）调研工具链（docs/export_support/）如何测量、续跑和合并结果。读完本文，你将掌握该页面的数据来源、每条状态的含义、如何重新运行调研刷新快照、如何注册新的被测用例，以及如何理解失败原因并判断一个 Kornia 算子能否导出到 ONNX 或跑通torch.compile。

该页面位于文档 Resources 目录下的get-started/export-support（入口见 docs/source/get-started/index.rst），对应源文件为docs/source/get-started/export-support.rst——注意这个.rst文件本身是由docs/generate_export_support.py在文档构建时根据快照 JSON 生成的，页面源码头部也明确写着"编辑 JSON 和生成脚本，而不是编辑本文件"。

一、这个页面解决什么问题

Kornia 是一个几何计算机视觉库，包含kornia.color、kornia.filters、kornia.geometry、kornia.augmentation、kornia.feature、kornia.losses、kornia.models等大量模块。用户在把模型部署到 ONNX Runtime、用torch.compile加速推理、或用torch.export做整图捕获时，最常遇到的问题是：某个算子到底能不能被捕获成一张静态图？失败原因往往是缺失的 ONNX lowering、数据相关的输出形状，或前向过程中的 Python 控制流——这些只有逐个算子实测才知道。

新增的迁移条目（changelog.d/+migration-022.added.md）记录了这一变化：

AnONNX, torch.compile and torch.export supportpage under Resources (get-started/export-support) lists every surveyed public operator, grouped by package and section, with its dynamo ONNX export,torch.exportandtorch.compileoutcome, the cause of each failure and a search box with result filters. It is rendered at build time from the committed snapshotdocs/source/_data/export_support.json;docs/export_support/run.pyreruns the survey to refresh it. (#4182)

也就是说，页面不是人工维护的表格，而是由一次**库级图捕获调研（survey）**驱动的：每个公开算子、模型和数据增强都用真实输入跑一遍 PyTorch 的三条图捕获路径，把结果提交为 JSON 快照，构建文档时再渲染成带搜索框和过滤器的交互式页面。

二、数据流：从调研到页面

完整链路如下：

1. 注册用例：在docs/export_support/cases_*.py中用case(...)登记每个被测算子（详见第四节）；
2. 探测：python docs/export_support/run.py并行启动两类 worker——harness.py做 ONNX 探测，probe_compile.py做torch.export/torch.compile探测——结果写入docs/export_support/results/（git-ignored）；
3. 合并：docs/export_support/merge.py校验所有结果文件与各自的inventory_*.json后，产出快照docs/source/_data/export_support.json（提交到仓库）；
4. 渲染：文档构建时由 docs/generate_export_support.py 读取该 JSON，生成docs/source/get-started/export-support.rst页面。

docs/export_support/README.md强调："nothing here runs during a docs build"——即调研脚本只在刷新快照时运行，构建页面时只有生成器读取已提交的 JSON。这样文档构建既不依赖网络、也不依赖可能很慢的探测过程，同时保证页面数据可审计、可复现。

页面渲染的关键参数（来自 generate_export_support.py 的状态表）如下：

每个状态还带有面向用户的简短原因（onnx_detail/export_detail/compile_detail），多个列共享同一原因时会合并显示，例如**ONNX, export, compile:** ...（见_details函数）。页面还额外统计"全部三条路径都通过"的配置数、被探测算子的去重数量（同一个算子的不同配置只算一个算子）等指标。

三、调研到底测量了什么

3.1 用例（case）的定义

一个case就是一个公开可调用对象（函数或nn.Module实例）加上一组具体的 CPU float32 张量输入，注册在某个cases_*.py文件中。根据README.md的划分：

例如 cases_misc.py 里定义了一批标准输入：IMG = torch.rand(1, 3, 32, 40)、GRAY = torch.rand(1, 1, 32, 40)、RGBA = torch.rand(1, 4, 32, 40)、VOL = torch.rand(1, 1, 8, 16, 16)等，然后通过add("RgbToGrayscale", KC.RgbToGrayscale(), IMG)这类调用批量注册。

3.2 三条探测路径

对每个用例，调研分别执行：

• ONNX（harness.py）：torch.onnx.export(dynamo=True, opset_version=18)，然后onnx.checker.check_model（full check），最后在 onnxruntime（CPU）中运行并与 eager 逐输出比较（默认atol=2e-4, rtol=1e-3，可用atol/rtol参数覆盖）；
• torch.export（probe_compile.py）：torch.export.export（默认非严格模式）；
• torch.compile（probe_compile.py#L153-L185）：先torch._dynamo.explain统计 graph break 次数，再用默认 inductor 后端、fullgraph=False跑torch.compile并与 eager 比较。

ONNX 路径还会记录图的操作类型集合（op_types）、导出图大小（size_kb）和耗时（time_s）；失败时记录根因错误（err_str会沿着__cause__/__context__链找最深层原因）以及涉及到的 kornia 源码帧（_kornia_frames，最多取最内层 6 帧），用于页面上显示kornia/xxx.py:行号定位信息。

3.3 随机算子的处理

含有随机性的算子无法逐值比较（图内 RNG 流与 eager 不同），因此check=False的用例在"导出并运行成功、输出形状正确且值有限"时记为ok-unverified。比如Normalize[random]这种把采样放进图内的配置，在快照里onnx与compile都是ok-unverified，并附注 "random operator: exported and ran; only output shape and finiteness are checked"（见 export_support.json 的实际记录）。

3.4 与 eager 的比较逻辑

比较函数close（harness.py）会先比对形状；布尔输出用逐元素相等；NaN 掩码必须一致且只比较非 NaN 部分；最后用np.allclose(atol, rtol)判断，失败时报告最大绝对差及超差元素数。输出展开用_flatten_outputs：支持Boxes、Keypoints、Se3、Quaternion、Image这类通过.data属性暴露张量的容器，也支持任意嵌套的 pytree 结构。

四、如何刷新快照：命令与机制

4.1 标准流程

docs/export_support/README.md给出了刷新快照的完整流程：

pixi run -e default install-docs && pixi run -e default build-docs # 先构建一次，让算子获得交叉引用 python docs/export_support/run.py # 8 核 CPU 约 2 小时，可续跑 pixi run -e default build-docs

第一步构建文档是为了生成 Sphinx inventory（docs/build/html/objects.inv），merge.py会用它把算子名转成页面上的交叉引用链接；如果没有该文件，所有算子将渲染为纯文本（merge.py 会打印警告但不会失败）。install-docs/build-docs这两个命令定义在 pixi.toml，其中build-docs实际执行uv run python -m sphinx -W -b html docs/source docs/build/html。

run.py支持的参数（见其 docstring 与main函数）：

python docs/export_support/run.py # 全部 6 个分组，约 2 小时（仅 CPU） python docs/export_support/run.py aug misc # 只跑指定分组 python docs/export_support/run.py --merge-only # 只基于已有 results 重新合并 JSON python docs/export_support/run.py --force # 即使结果不完整也强制合并

分组名取自GROUPS = ["aug", "feature", "geomA", "geomB", "misc", "models"]（run.py），传了未知分组会直接报错退出。

4.2 并行与续跑机制

run.py用ThreadPoolExecutor为每个分组启动两个进程（onnx:<group>跑cases_<group>.py，compile:<group>跑probe_compile.py），即最多 12 个并发 worker。每个 case 在运行前会先写入一个crashed标记，这样进程若因 segfault/OOM 硬崩溃，重启后该 case 的失败状态是可见的。

续跑（resume）的核心在 harness.py 的 prepare_resume：每个探测会写一个inventory_*.json，记录本次运行的 kornia revision（git rev-parse --short HEAD，含-dirty后缀）和 torch 版本。已存在的记录只有在 inventory 的(revision, torch)与当前完全一致时才会被复用，否则全部丢弃重新测量——不同版本下测得的结果不会被续用，避免跨版本混用数据。

probe_compile.py还引入了"中毒"（poisoned）机制：如果某个 case 让 dispatcher 状态损坏（PythonDispatcherTLS not set），worker 以退出码POISON_EXIT = 3退出，run.py会重启它（最多MAX_RESTARTS = 50次），由于结果可续跑，重启后会自动跳过已完成 case 继续剩余工作。

4.3 合并的保护措施

merge.py的_validate会做三重校验，任何一项不满足都会拒绝写快照（除非--force）：

1. 每个结果文件旁边必须有对应的inventory_*.json；
2. 每个 case 必须有非 crash 标记的真实记录（否则报告缺多少个，如12 of 40 cases have no record, e.g. [...]）；
3. 所有分组的 inventory 必须声明相同的 kornia revision 和 torch 版本。

这样部分运行永远不会静默地少算覆盖率。合并时还排除了 crash 标记和由中毒解释器导致的假 eager 失败记录（_load_results中的bogus判断）。合并出的 JSON 头部记录了生成日期、kornia/torch/onnx/onnxruntime/onnxscript 版本、Python 版本、opset（18）、设备（cpu），并按 pre-commit 的pretty-format-json期望格式（indent=2, sort_keys=True）写出，保证重新生成时 diff 干净。

五、失败原因的自动归类

页面 Details 列的每个失败原因都由merge.py的_reason函数自动生成，分为几类：

• 缺失 ONNX lowering：错误信息匹配No ONNX function found for <OpOverload(op='aten.xxx'>时，用ATEN_NAMES表映射成用户叫法，例如_linalg_svd → torch.linalg.svd、linalg_qr → torch.linalg.qr、histc → torch.histc、searchsorted → torch.searchsorted；
• 数据相关（data-dependent）：_CAUSES表按错误子串匹配，例如：
  • "Could not extract specialized integer from>from harness import case, run_cases CASES = [case("rgb_to_grayscale", "color", K.color.rgb_to_grayscale, [torch.rand(1, 3, 32, 32)]), ...] if __name__ == "__main__": run_cases(CASES, sys.argv[1], only=sys.argv[2:] or None)

    run_cases会先检查重复 case 名（重复则拒绝运行），并且对每个 case 先写 crash 标记再运行、运行后立即落盘，保证中断/崩溃后重启能无缝续跑；每个 case 也会打印一行name status size time的实时进度，最后输出按状态计数的 SUMMARY。only参数支持按 case 名、分组名或名字前缀过滤。

    七、实测数据的样子与版本前提

    以提交的快照 docs/source/_data/export_support.json（约 2.2 万行）为例，前两条记录分别是Normalize（三条路径全部ok，0 graph breaks）和Normalize[random]（onnx/compile为ok-unverified，附随机算子说明）——这正是文档里所说"增强类算子探测两次：参数作为图输入（确定性图）与采样放进图内（[random]）"的体现。

    页面头部还会展示三项卡片指标：ONNX 导出、torch.export、torch.compile各自的supported / probed比例，并注明快照的测量环境（kornia、torch、onnx、onnxruntime、onnxscript 版本及 Python 版本）。这些数字与版本强绑定——README 明确提醒：除非本次重跑的目的就是升级版本，否则请用与快照头部一致的 torch / onnx / onnxruntime / onnxscript 版本运行探测，否则合并出的快照会与页面环境不一致。页面同样建议读者查阅docs/source/get-started/onnx.rst了解如何串联多个导出的图，查阅docs/source/get-started/gpu-acceleration.rst了解torch.compile的实际用法。

    八、小结

    这个新增的 Export Support 页面把"Kornia 算子能否被三种图捕获路径支持"从人工经验变成了可审计、可续跑、可复现的工程事实：cases_*.py定义测量范围，harness.py/probe_compile.py执行三类探测，run.py编排并行与续跑，merge.py用 inventory 校验保证合并安全，generate_export_support.py在构建时把提交的快照渲染成带搜索框、过滤器和逐条失败原因的可交互文档。对开发者而言，它的价值是双重的：既是一张部署前查证算子兼容性的权威表格，也是一套可以自行扩展和刷新的测量流水线——新增一个公开算子，注册一个case(...)，重跑run.py并提交新快照即可。

    【免费下载链接】kornia🐍 Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/GitHub_Trending/ko/kornia