Kornia 相机标定 tilt 检查兼容 torch.compile(fullgraph=True) 的修复详解

Kornia 相机标定 tilt 检查兼容 torch.compile(fullgraphTrue) 的修复详解【免费下载链接】kornia Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/gh_mirrors/ko/kornia导读本文围绕 Kornia 合入的修复对应 changelog 片段PR #4391展开distort_points与undistort_points原本在 host 端读取张量形式的 tilt 畸变系数τx、τy来决定是否走 tilt 投影分支这种由张量值驱动控制流的写法在torch.compile(fullgraphTrue)图捕获阶段无法读取系数数值会破坏编译。修复后编译/导出路径无条件应用 tilt 分支零角度时退化为单位矩阵同时完整保留 eager 与 TorchScript 行为。读完本文你将理解该 bug 的根因、Kornia 采用的兼容策略、底层tilt_projection的数学实现以及对应的编译回归测试如何验证正确性。修复背景为什么 tilt 检查会破坏 fullgraph 编译Kornia 的镜头畸变模型在径向k1..k6、切向p1,p2、薄棱镜s1..s4之外还支持 tilt 畸变对应dist系数向量的第 13、14 个元素τx, τy。distort_points与undistort_points接受 4/5/8/12/14 长度的系数向量并在内部零填充到 14 个元素见 distort.py 与 undistort.py。修复前的代码用如下条件判断是否应用 tilt 变换if torch.any(dist[..., 12] ! 0) or torch.any(dist[..., 13] ! 0): tilt tilt_projection(dist[..., 12:13], dist[..., 13:14]) ...问题在于torch.any(dist[..., 12] ! 0)是对张量值的运行时判断。在torch.compile(fullgraphTrue)图捕获阶段Dynamo 无法读取也不允许读取GPU/CPU 张量的具体数值来决定控制流这种 tensor-dependent 的分支要么导致图捕获失败要么迫使 Dynamo 将系数值特化specialize进编译结果——一旦运行时传入不同的 tilt 系数编译产物就会给出错误结果。这正是 test_compile.py 注释中所记录的 #4286 问题tensor-dependent tilt checks must not break capture or specialize on coefficient values。修复方案捕获期检测 无条件应用 tilt 分支修复思路非常直接图捕获阶段无法看系数值那就干脆不看——在编译或导出时无条件执行 tilt 投影分支由于τx τy 0时 tilt 投影矩阵恒等于单位矩阵eye(3)无条件分支不会改变数学结果零 tilt 的投影就是恒等变换。改动集中在两个文件的同一段逻辑两处实现刻意保持一致distort.py前向畸变# Graph capture cannot read the coefficient values on the host. Apply the tilt unconditionally # while compiling or exporting; zero angles give the identity. Keep eager and scripted behavior. capture is_exporting() if not torch.jit.is_scripting(): capture capture or is_compiling() if capture or torch.any(dist[..., 12] ! 0) or torch.any(dist[..., 13] ! 0): tilt tilt_projection(dist[..., 12:13], dist[..., 13:14]) points_untilt torch.stack([xd, yd, torch.ones_like(xd)], -1) tilt.transpose(-2, -1) xd points_untilt[..., 0] / points_untilt[..., 2] yd points_untilt[..., 1] / points_untilt[..., 2]undistort.py逆畸变capture is_exporting() if not torch.jit.is_scripting(): capture capture or is_compiling() if capture or torch.any(dist[..., 12] ! 0) or torch.any(dist[..., 13] ! 0): inv_tilt tilt_projection(dist[..., 12:13], dist[..., 13:14], True) if inv_tilt.dim() 2: inv_tilt inv_tilt.unsqueeze(0) x, y transform_points(inv_tilt, torch.stack([x, y], dim-1)).unbind(-1)要点拆解分支条件变为capture or 数值判断只要处于图捕获环境就进入 tilt 分支不再依赖张量值eager 模式下仍保留原来的数值判断零 tilt 时跳过多余计算。is_exporting()与is_compiling()的语义二者定义在 kornia/core/utils.py。is_compiling()判断是否处于torch.compile或torch.export捕获中内部解析torch.compiler.is_compiling缺失时回退到torch._dynamo.is_compilingis_exporting()判断是否处于torch.export或 dynamo ONNX 导出器捕获中并且torch.jit.unused装饰保证 TorchScript 下恒为False因此该 guard 可安全用于脚本化函数。torch.jit.is_scripting()保护显式排除 TorchScript 路径保持eager 和 scripted 行为与修复前完全一致——这正是 changelog 中所说的preserving eager behavior。底层支撑tilt_projection 的数学实现tilt 分支调用的tilt_projectiondistort.py构造 OpenCV 风格的 tilt 投影矩阵约定R Ry(tauy) Rx(taux)并支持正逆两个分支return_inverseFalse返回Pz R与 OpenCV 的 tilt projection 一致用于distort_points的加 tilt方向return_inverseTrue返回Pz R的逆矩阵R.T invPz即undistort_points使用的去 tilt方向可在仓库自带的参考值上复现 OpenCV 的undistortPoints。两个分支在τx τy 0时都精确返回eye(3)旋转矩阵为单位阵、R[..., 2, 2] 1这是无条件应用 tilt 分支也不会改变结果的数学保证。相关约定在测试中有专门覆盖见 test_distort.py零角度恒等性、正逆分支互为逆矩阵、前向分支等于Pz R等断言对应问题 #4276。回归测试fullgraph 编译下的逐值一致性验证本次修复配套了专门的编译回归测试 test_compile.pypytest.mark.parametrize(op, [distort_points, undistort_points]) pytest.mark.parametrize(num_coefficients, [4, 14]) def test_dynamo_fullgraph(self, device, dtype, torch_optimizer, op, num_coefficients): counter CompileCounter() counted torch.compile(op, backendcounter, fullgraphTrue) compiled torch_optimizer(op, fullgraphTrue) for tilt in [0.0, 0.1]: if num_coefficients 14: dist[:, 12] tilt dist[:, 13] -2 * tilt expected op(points, camera, dist) self.assert_close(counted(points, camera, dist), expected) self.assert_close(compiled(points, camera, dist), expected) assert counter.frame_count 1测试覆盖了两个关键维度op × 系数长度distort_points与undistort_points各自在 4 系数无 tilt 槽位和 14 系数含 tilt 槽位下编译tilt 取值分别以0.0零 tilt编译分支无条件走 tilt但数学上应为恒等和0.1真实 tiltτy -2τx验证。断言要求fullgraphTrue编译后的输出与 eager 参考值assert_close一致同时CompileCounter记录frame_count 1——即整个函数被成功捕获进单个编译图没有因 host 端取值而触发图重建或回退。修复前这类测试会因无法在捕获期读取dist值而失败。此外既有功能测试也验证了修复不改变 eager 语义例如 test_undistort.py 的test_convention_distort_undistort_round_trip_with_tilt_4276在非零径向/切向系数叠加非对称 tilt[0.1, 0.2]时undistort_points(distort_points(...))往返仍能闭合到参考点。实战指引如何在你的代码中编译这两个算子修复生效后可以放心地对标定流程做端到端图编译import torch from kornia.geometry.calibration import distort_points, undistort_points points torch.tensor([[[1.0, 2.0], [6.0, 5.0]]]) K torch.tensor([[[100.0, 0.0, 4.0], [0.0, 100.0, 3.0], [0.0, 0.0, 1.0]]]) dist torch.zeros(1, 14) dist[:, :4] torch.tensor([0.1, 0.01, 0.001, 0.0001]) dist[:, 12] 0.1 # taux dist[:, 13] -0.2 # tauy compiled_distort torch.compile(distort_points, fullgraphTrue) compiled_undistort torch.compile(undistort_points, fullgraphTrue) d compiled_distort(points, K, dist) u compiled_undistort(d, K, dist)注意事项dist系数向量长度必须为 4、5、8、12 或 14 之一其他长度会抛出ValueError短向量内部自动零填充到 14因此 4 系数与 14 系数零填充版本结果一致见 distort.py。undistort_points的逆变换是num_iters5的固定点迭代而非闭式解往返精度受迭代次数与 dtype 限制float16在默认迭代次数下可能已触及舍入下限float32/float64增大num_iters可继续改善见 undistort.py。若同时使用 ONNX 导出is_exporting()会走与编译相同的无条件 tilt路径零 tilt 时投影为恒等导出图行为一致。该修复依赖 Kornia 的捕获检测工具位于 kornia/core/utils.py其同时被其他 export-safe 路径如闭式逆、sort类中位数等复用。小结本次修复的实质是把依赖张量值的数据依赖分支改写为环境感知的条件分支图捕获期间无条件走数学上等价零角度恒等的 tilt 分支eager 与 TorchScript 则维持原有数值判断。它同时保证了三点——fullgraphTrue可成功捕获单图、编译/导出产物与 eager 输出逐值一致、以及零 tilt 时的恒等语义。若要深入源码可从 distort.py 与 undistort.py 的 tilt 分支入手配合 test_compile.py 与 test_distort.py 中的约定测试对照阅读。【免费下载链接】kornia Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/gh_mirrors/ko/kornia创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考