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 快照构建文档时再渲染成带搜索框和过滤器的交互式页面。二、数据流从调研到页面完整链路如下注册用例在docs/export_support/cases_*.py中用case(...)登记每个被测算子详见第四节探测python docs/export_support/run.py并行启动两类 worker——harness.py做 ONNX 探测probe_compile.py做torch.export/torch.compile探测——结果写入docs/export_support/results/git-ignored合并docs/export_support/merge.py校验所有结果文件与各自的inventory_*.json后产出快照docs/source/_data/export_support.json提交到仓库渲染文档构建时由 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 的状态表如下状态单元格显示含义okyes成功捕获且结果与 eager 执行一致ok-breaksyes, graph breaks编译成功且与 eager 一致但 dynamo 在若干处 graph break 回退到 Pythonok-unverifiedyes (random)捕获并运行成功因含随机算子只比较输出形状与有限性mismatchmismatch捕获并运行成功但结果与 eager 不同failno捕获或转换失败Details 列说明原因n/an/a不是张量图文件 I/O、枚举、Python 容器、stub或无法用探测输入运行每个状态还带有面向用户的简短原因onnx_detail/export_detail/compile_detail多个列共享同一原因时会合并显示例如**ONNX, export, compile:** ...见_details函数。页面还额外统计全部三条路径都通过的配置数、被探测算子的去重数量同一个算子的不同配置只算一个算子等指标。三、调研到底测量了什么3.1 用例case的定义一个case就是一个公开可调用对象函数或nn.Module实例加上一组具体的 CPU float32 张量输入注册在某个cases_*.py文件中。根据README.md的划分文件覆盖范围cases_aug.pykornia.augmentation2D、3D、容器、自动策略cases_feature.pykornia.feature检测器、描述子、匹配器cases_geomA.pykornia.geometry的变换、转换、相机、标定cases_geomB.pykornia.geometry其余部分对极几何、李群、Boxes 等cases_misc.pykornia.filters、color、enhance、morphology、contrib、losses、metrics、utils、sensors、io、nerf、tracking、xcases_models.pykornia.models权重未缓存时会自动下载例如 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 三条探测路径对每个用例调研分别执行ONNXharness.pytorch.onnx.export(dynamoTrue, opset_version18)然后onnx.checker.check_modelfull check最后在 onnxruntimeCPU中运行并与 eager 逐输出比较默认atol2e-4, rtol1e-3可用atol/rtol参数覆盖torch.exportprobe_compile.pytorch.export.export默认非严格模式torch.compileprobe_compile.py#L153-L185先torch._dynamo.explain统计 graph break 次数再用默认 inductor 后端、fullgraphFalse跑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 不同因此checkFalse的用例在导出并运行成功、输出形状正确且值有限时记为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 的比较逻辑比较函数closeharness.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 inventorydocs/build/html/objects.invmerge.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.pycompile:group跑probe_compile.py即最多 12 个并发 worker。每个 case 在运行前会先写入一个crashed标记这样进程若因 segfault/OOM 硬崩溃重启后该 case 的失败状态是可见的。续跑resume的核心在 harness.py 的 prepare_resume每个探测会写一个inventory_*.json记录本次运行的 kornia revisiongit rev-parse --short HEAD含-dirty后缀和 torch 版本。已存在的记录只有在 inventory 的(revision, torch)与当前完全一致时才会被复用否则全部丢弃重新测量——不同版本下测得的结果不会被续用避免跨版本混用数据。probe_compile.py还引入了中毒poisoned机制如果某个 case 让 dispatcher 状态损坏PythonDispatcherTLS not setworker 以退出码POISON_EXIT 3退出run.py会重启它最多MAX_RESTARTS 50次由于结果可续跑重启后会自动跳过已完成 case 继续剩余工作。4.3 合并的保护措施merge.py的_validate会做三重校验任何一项不满足都会拒绝写快照除非--force每个结果文件旁边必须有对应的inventory_*.json每个 case 必须有非 crash 标记的真实记录否则报告缺多少个如12 of 40 cases have no record, e.g. [...]所有分组的 inventory 必须声明相同的 kornia revision 和 torch 版本。这样部分运行永远不会静默地少算覆盖率。合并时还排除了 crash 标记和由中毒解释器导致的假 eager 失败记录_load_results中的bogus判断。合并出的 JSON 头部记录了生成日期、kornia/torch/onnx/onnxruntime/onnxscript 版本、Python 版本、opset18、设备cpu并按 pre-commit 的pretty-format-json期望格式indent2, sort_keysTrue写出保证重新生成时 diff 干净。五、失败原因的自动归类页面 Details 列的每个失败原因都由merge.py的_reason函数自动生成分为几类缺失 ONNX lowering错误信息匹配No ONNX function found for OpOverload(opaten.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], onlysys.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三条路径全部ok0 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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考