Kornia 精度修复ConvQuadInterp3d/conv_quad_interp3d的 float16 梯度溢出NaN问题与 float32 提升方案【免费下载链接】kornia Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/GitHub_Trending/ko/kornia导读本文基于 Kornia 仓库的变更记录changelog.d/migration-077.fixed.md深入剖析kornia.geometry.subpix中三维亚像素定位算子ConvQuadInterp3d/conv_quad_interp3d在float16半精度输入下产生 NaN 梯度的根因以及官方采用的修复方案将克拉默Cramer规则求解 Hessian 行列式的过程提升promote到float32完成再在最后将位移结果转回float16。文章同时结合源码实现kornia/geometry/subpix/spatial_soft_argmax.py与回归测试tests/geometry/subpix/test_spatial_softargmax.py展开说明帮助读者理解半精度数值下溢的机理、修复的前向/反向行为以及为何bfloat16和float32不受影响。背景3D 尺度空间极值的亚像素定位conv_quad_interp3d是 Kornia 在kornia/geometry/subpix模块中提供的函数另有模块封装类ConvQuadInterp3d用于对 3D 尺度空间响应体中的局部极大值做亚像素定位。它的输入是一个形状为(B, C, D, H, W)的响应金字塔response pyramid其中D为尺度scale维度H、W为空间维度。其核心思路是通过nms3d找到局部极大值候选在每个极大值周围提取3 × 3 × 3邻域拟合一个三维二次函数quadratic求解该二次函数的极值点得到亚体素sub-voxel位移(sx, sy, ss)若某轴位移超过max_subpixel_shift默认 0.6则将整数中心沿该轴移动一步并重新求解最多迭代n_iters默认 5次。该函数的文档字符串给出了一个直观示例spatial_soft_argmax.py input torch.randn(2, 3, 5, 64, 64) coords, vals conv_quad_interp3d(input, n_iters5) coords.shape torch.Size([2, 3, 3, 5, 64, 64]) vals.shape torch.Size([2, 3, 5, 64, 64])输出的coords_max在每个 NMS 极大值位置存放细化的[scale, x, y]亚像素坐标y_max存放经二次函数修正后的响应值。该算子在 SIFT 等特征检测器的尺度空间极值精化中被广泛使用Kornia 的 SIFT 实现位于 kornia/feature/sift/scale_space.py其中同样调用_solve_cramer_sym3x3。问题根因Hessian 行列式在 float16 下的梯度下溢求解的核心_solve_cramer_sym3x3无论是conv_quad_interp3d还是同模块的迭代版本iterative_quad_interp3d最终都要通过对称 3×3 线性系统H · [sx, sy, ss]ᵀ [r0, r1, r2]ᵀ求位移其中 Hessian 矩阵H [[dxx, dxy, dxs], [dxy, dyy, dys], [dxs, dys, dss]]的元素来自 3×3×3 邻域的有限差分二阶导数。该系统的求解由内部函数_solve_cramer_sym3x3承担spatial_soft_argmax.py其 docstring 明确描述了 Cramer 法则的求解语义对于行列式绝对值小于eps默认1e-7的奇异/近奇异系统返回的位移在数值上没有意义调用方通过布尔掩码solved将其丢弃。数值量级的数学分析问题出在行列式det的数值量级上。如变更记录所述Hessian 行列式是三个二阶导数的乘积。对取值在[0, 1]的响应图而言二阶差分如dxx p_xp - 2*c000 p_xm的量级约为1e-2甚至更小三者的乘积因此落在1e-4及以下。这个量级本身仍然大于eps 1e-7所以solved det.abs() eps判定成立——前向传播没有问题前向只需除以一次safe_det结果保持有限。反向传播的1 / det²爆炸真正的杀手在反向传播。源码中对位移的求解形如sx (r0 * cf00 - ...) / safe_det而safe_det torch.where(solved, det, torch.ones_like(det))。PyTorch 对num / safe_det求梯度时会按链式法则乘上1 / safe_det**2因子。也就是说反向传播的梯度幅度与行列式的平方成反比。当det ≈ 1e-4时1 / det² ≈ 1e8而当det更小比如1e-51 / det²高达1e10。float16半精度能表示的最小正规数finfo(float16).tiny约为6.1e-5指数范围远窄于float32。于是梯度中1 / det²这样的量在float16下直接下溢为 0 或上溢为 inf在响应体的一部分体素上梯度溢出后经归约reduce变成NaN。这正是变更记录所指出的现象the backward scales by1 / det**2, whichfloat16cannot represent ... so the gradient overflowed and reduced to NaN over part of the volume。修复方案求解过程提升到 float32位移结果再转回源码中的实现证据修复落点在_solve_cramer_sym3x3内部spatial_soft_argmax.py源码注释完整复述了上述分析# float16 cannot carry this solve. The determinant is a product of three # second derivatives, so for a [0, 1] response it lands around 1e-4 and # below — still above eps, so solved admits it. The forward divides # by it once and stays finite, but the backward of num / safe_det # scales by 1 / safe_det**2, and that square is not representable in # float16 (finfo.tiny is 6.1e-5), so the gradient becomes inf and reduces # to NaN. bfloat16 keeps float32s exponent range and is unaffected.实现逻辑非常直接记录输入 dtype 后仅在float16时将全部输入提升到float32完成 Cramer 求解后再把位移结果sx, sy, ss转回float16in_dtype dxx.dtype if in_dtype torch.float16: dxx, dyy, dss dxx.float(), dyy.float(), dss.float() dxy, dxs, dys dxy.float(), dxs.float(), dys.float() r0, r1, r2 r0.float(), r1.float(), r2.float() cf00 dyy * dss - dys * dys # cofactor M00 ... sx (r0 * cf00 - ...) / safe_det ... if in_dtype torch.float16: sx, sy, ss sx.to(in_dtype), sy.to(in_dtype), ss.to(in_dtype)由于中间运算全部在float32完成1 / det²不再触碰float16的窄指数范围反向传播自然保持有限值。与 #4231 修复的一致性变更记录提到这一做法与先前合并的 #4231changelog.d/migration-020.breaking.md保持一致——后者修复了create_meshgrid/create_meshgrid3d在float16/bfloat16下归一化坐标不一致的问题其方案同样是半精度网格在float32中构建与归一化、仅最后收窄一次。可见 Kornia 在涉及数值敏感中间量的半精度路径上采用了统一的提升-计算-收窄模式。为什么 bfloat16 不受影响bfloat16虽然尾数位更少7 位但指数范围与float32完全一致8 位指数因此能够承载1e-4乃至更小的行列式平方倒数不会像float165 位指数那样溢出。这正是变更记录中bfloat16keepsfloat32s exponent range and was never affected的含义也是修复没有对bfloat16做提升的原因。行为变更与兼容性说明输入 dtype修复前修复后float16部分体素反向梯度为 NaNbug求解提升至float32梯度保持有限位移结果转回float16bfloat16正常输出不变指数范围与float32相同未受影响float32正常输出不变float64正常输出不变走 CUDA 打包路径见下需要注意的是bfloat16与float32的前向输出在修复前后完全一致——提升只发生在float16分支且不改变有效数字计算结果float16输入在修复前的前向本身是有限的输出数值也没有被重排。CUDA 路径的对应处理在 CUDA 上_solve_cramer_sym3x3对float32/float64的一维输入会转交_solve_cramer_sym3x3_cudaspatial_soft_argmax.py——该函数将四个 Cramer 行列式打包成一次张量运算以换取更少的 kernel 启动而其 docstring 特别注明the scalar implementation is faster on CPU and retains the float16 solves promotion behavior即 CUDA 打包路径只服务float32/float64float16的提升行为由标量路径统一承担两条路径各司其职。回归测试验证仓库为该修复提供了专门的回归测试test_float16_backward_is_finitetests/geometry/subpix/test_spatial_softargmax.pydef test_float16_backward_is_finite(self, device): # The Hessian determinant is a product of three second derivatives, so for a # [0, 1] response it lands around 1e-4 and below. The forward divides by it # once and stays finite, but the backward scales by 1 / det**2, which float16 # cannot represent, so the gradient used to be NaN over part of the volume. # Pinned in float16 specifically: bfloat16 keeps float32s exponent range. torch.manual_seed(0) sample torch.rand(1, 1, 5, 40, 40, devicedevice, dtypetorch.float16, requires_gradTrue) coords, vals kornia.geometry.subpix.ConvQuadInterp3d(strict_maxima_bonus0.0)(sample) grad torch.autograd.grad(coords.sum() vals.sum(), sample)[0] assert torch.isfinite(grad).all()测试要点构造(1, 1, 5, 40, 40)的float16随机输入并打开requires_grad用strict_maxima_bonus0.0关闭严格极大值加分隔离纯插值路径对输出坐标和响应值求和后反向求梯度断言梯度处处有限torch.isfinite(grad).all()测试注释特别强调仅在float16固定该行为因为bfloat16保持float32的指数范围、本就不会触发该 bug。同文件中的其他测试如test_peak_at_center、test_subpixel_shift、test_dynamo等则持续守护前向定位精度与torch.compile兼容性与本次数值修复形成互补。对使用者的影响与建议训练/可微流程如果你的 pipeline 以float16如 AMP 混合精度运行ConvQuadInterp3d或conv_quad_interp3d修复前可能观测到定位梯度中的 NaN 与训练不稳定升级到包含该修复的 Kornia 版本后梯度保持有限无需改代码。前向推理float32与bfloat16的输出与修复前逐位一致不需要更新推理结果或基线float16的前向本身有限主要变化体现在反向。数值直觉对于响应值在[0, 1]量级的尺度空间二阶差分乘积天然落在1e-4量级任何除以行列式的算子都应警惕半精度下的平方倒数爆炸提升-计算-收窄是 Kornia 官方验证过的通用解法参见 #4231 的网格修复。相关参考核心实现位于 kornia/geometry/subpix/spatial_soft_argmax.pySIFT 尺度空间中的同类求解位于 kornia/feature/sift/scale_space.py半精度相关回归测试位于 tests/geometry/subpix/test_spatial_softargmax.py。【免费下载链接】kornia Geometric Computer Vision Library for Spatial AI项目地址: https://gitcode.com/GitHub_Trending/ko/kornia创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
