Kornia 修复 create_meshgrid 半精度归一化坐标在 eager 与 torch.compile 下的不一致

• 计算机视觉
• 深度学习
• 图像处理

导读

本文深入解析 Kornia 一次针对数值一致性的重要修复（issue #4030，PR #4231）：create_meshgrid与create_meshgrid3d生成的float16/bfloat16半精度归一化坐标网格，在 eager 执行与torch.compile编译执行下原本会得到不同的数值，导致模型一旦被编译包裹，采样网格就悄然发生漂移。文章将说明问题的根源（多次舍入 vs 单次舍入、窄化后拓宽被编译器折叠）、修复方案（统一在float32中构建与归一化、仅做一次窄化）、对HomographyWarper等下游算子的影响，以及剩余的 CUDA 平台边界。读完你将理解半精度网格坐标的舍入机制，并能在自己的训练/推理代码中规避同类陷阱。

一、问题背景：归一化坐标网格的两种生成路径

create_meshgrid是 Kornia 几何变换的核心工具，位于 kornia/geometry/grid.py，它生成形状为(1, H, W, 2)的坐标网格。当normalized_coordinates=True时，网格被归一化到[-1, 1]区间，与 PyTorchtorch.nn.functional.grid_sample的约定保持一致。归一化公式为：

x_norm = (xs / (size - 1) - 0.5) * 2

其中xs是由torch.arange生成的像素索引 ramp（0, 1, 2, ..., size-1）。

在本次修复之前，create_meshgrid内部存在两条不同的执行路径：

1. eager 路径：(xs / (size - 1) - 0.5) * 2被当作三次独立的半精度运算，即除法、减法、乘法各自在float16/bfloat16下执行，每步都做一次舍入；
2. graph-capture 路径（torch.compile/torch.jit.trace捕获的图）：编译器（Inductor）把整条表达式链提升到float32计算、只在结果存储时舍入一次。

与此同时，代码在进入 capture 分支前会先把半精度 ramp 物化出来，capture 分支再将其拓宽为工作精度——这种"先窄化再拓宽"（narrow-then-widen）的往返，恰好被编译器视为冗余而折叠掉：Inductor 直接保留了精确的整数索引，而 eager 路径则是用已经舍入过的半精度分子去做除法。

这两点叠加的后果是：当网格宽度超过2**p（p 为有效数位数，bfloat16为 8 位、float16为 11 位）时，两条路径除法的分子已经不同——即 ramp 整数在超过2**p后无法被半精度精确表示，eager 的半精度分子带上了舍入误差，而编译路径的精确分子没有。

二、量化影响：模型被 torch.compile 包裹后网格悄然漂移

由于上述机制，一个半精度网格在模型被torch.compile包裹的瞬间就会改变数值。修复说明中给出的实测量级（bfloat16，[-1, 1]归一化范围）：

0.0117 恰好是 bfloat16 在[-1, 1]区间上约 1 个 ulp（2**-7 ≈ 0.0078至2**-6 ≈ 0.0156之间）的数值；宽度 6000 时偏差约 0.0156，接近 2 个 ulp。对精度敏感的采样任务（如特征点对齐、光流、深度图 warp）而言，这种因"是否编译"而产生的不可复现差异是难以排查的。

值得强调的是，差异并非编译器的 bug，而是两条路径对同一数学表达式的舍入策略不同：

• eager：半精度逐步运算，逐步舍入（3 次舍入）；
• compiled：float32整链运算，仅存储时舍入（1 次舍入）；
• eager 的窄化-拓宽往返制造了带误差的分子，编译路径将其折叠并保留精确索引。

三、修复方案：float32 工作精度 + 单次窄化

修复的核心思路在 kornia/geometry/grid.py 中清晰可见，实现为三步：

ramp_dtype = dtype if dtype is not None else torch.empty(0, device=device).dtype # 只有半精度且需要归一化时才拓宽 widened = normalized_coordinates and ramp_dtype in (torch.float16, torch.bfloat16) work_dtype = torch.float32 if widened else ramp_dtype xs = torch.arange(width, device=device, dtype=work_dtype) ys = torch.arange(height, device=device, dtype=work_dtype)

1. 统一工作精度：当输入 dtype 是float16/bfloat16且normalized_coordinates=True时，把work_dtype提升为float32，ramp 直接以float32生成——eager 与 capture 分支看到的是同一份分子；
2. 统一归一化算术：无论 eager 还是 capture 分支，都在float32下完成(xs / (size - 1) - 0.5) * 2（或align_corners=False的半像素映射(2 * xs + 1) / width - 1），算术过程在所有路径上完全一致；
3. 只窄化一次：归一化完成后才在末尾做单次窄化回目标 dtype：

if widened: xs = xs.to(ramp_dtype) ys = ys.to(ramp_dtype)

这样，eager、traced 与 compiled 三种执行模式得到完全一致的网格，且 eager 数值移动到了"正确舍入"的一侧——这正是本变更属于breaking change的原因：eager 路径下半精度网格的数值也发生了改变（向正确舍入值靠拢）。create_meshgrid3d采用了完全相同的策略，见 kornia/geometry/grid.py。

3.1 配套修改：arange 取代 linspace 生成 ramp

本次修复还顺带把 ramp 的生成从linspace改为arange（kornia/geometry/grid.py 的注释说明了理由）：

• linspace会先把端点舍入到坐标 dtype，再从舍入后的端点反向填充 ramp 上半部分，因此只有当size - 1是可精确表示的整数（size <= 2**p + 1，bfloat16 为 257、float16 为 2049）时步长才精确为 1，超出后每个索引会偏离正确舍入的列索引最多 1 ulp；而 Inductor 将同一个调用降低为逐索引计算、正确舍入，导致 eager 与 compiled 的 ramp 相差最多 1 ulp；
• arange同时能在导出（export）场景下保留动态输出长度，linspace会对符号steps做特化。

此外，默认 dtype 通过torch.empty(0, device=device).dtype读取，而不是torch.get_default_dtype()：因为 TorchScript 没有后者的内建支持，且对无输入的工厂调用，TorchScript 优化执行器会在首次运行后做常量折叠，若直接读取默认 dtype 会把 ramp 冻结在首次调用时的默认值上。带上device的empty张量使读取依赖输入参数，从而避免折叠。

3.2 位级不变性保证

由于float32/float64和整数网格本来就不会被拓宽，本次修复对它们位级完全不变（bit-for-bit unchanged）：

• float32、float64网格：widened恒为False，work_dtype == ramp_dtype，执行路径与修复前一致；
• 整数网格与像素 ramp（normalized_coordinates=False）：不做任何归一化运算，不受影响；
• align_corners两种约定（True的像素中心映射与False的半像素边缘映射）的公式在float32中都是精确执行的，只有最终一次窄化落在目标 dtype 上。

四、HomographyWarper：预计算网格与即时网格的一致性

HomographyWarper（kornia/geometry/transform/homography_warper.py）是网格问题的直接下游受益者。它内部存在两条采样网格来源：

1. 预计算网格：__init__中通过create_meshgrid(height, width, normalized_coordinates=..., align_corners=...)构建基础网格（homography_warper.py），随后precompute_warp_grid(src_homo_dst)调用warp_grid（kornia/geometry/transform/imgwarp.py）用transform_points变换得到缓存的_warped_grid；
2. 即时网格：forward收到src_homo_dst时，走函数式homography_warp（kornia/geometry/transform/imgwarp.py），在调用时按输入 dtype 现算网格。

修复前，HomographyWarper的预计算网格在float32下构建、而即时路径的网格按输入 dtype（半精度）构建，两条路径的半精度数值不同，导致同一个 warp 在"预计算"与"直接传入 homography"两种用法下结果不一致。修复后两者共享create_meshgrid的 float32 工作精度逻辑，预计算与即时路径在半精度下完全一致，顺带修复了：

• test_homography_warper：在float16与bfloat16两个 dtype 下均通过；
• test_translation[shape1]：在float16下通过。

五、残余边界与测试基线变化

修复并非声称半精度在所有平台上达到位级一致，文档明确指出了剩余边界：

• CUDA + float16：最多仍有 1 个 ulp（坐标 dtype 的 ulp）差异，因为 ATen 与 Triton 的float32除法并非位级相同，除以size - 1时产生不同的舍入；
• CUDA + bfloat16：bfloat16 有效位数少，能吸收这一 ulp 级差异；
• CPU 内核：eager 与编译路径完全一致（bit-exact），Linux CPU 上的半精度基线测试随之更新。

5.1 测试容忍度分析：test_translation 的 atol=1e-4

测试 tests/geometry/transform/test_homography_warper.py 的test_translation以atol=1e-4断言平移后的像素对齐。修复后该测试在bfloat16下新增了两个待记录（xfail）条目：

• 此前已有[shape1]、[shape3]在bfloat16下无法达到atol=1e-4；
• 现在[shape0]、[shape2]也加入其中，即bfloat16 下四种 shape 全部无法满足atol=1e-4。

修复说明给出了一个反直觉的原因：旧网格之所以有两个 shape "碰巧通过"，并不是因为旧网格更精确，而是因为旧网格更大的误差，在经grid_sample的 bfloat16 反归一化（unnormalize）后恰好往返为一个精确的像素索引——即错误的网格与错误的采样舍入相互抵消。修复后网格数值移到正确舍入一侧，这种巧合消失，测试反而暴露了 bfloat16 本身的精度天花板：在[-1, 1]区间上 bfloat16 的 ulp 约为 0.0078，远超atol=1e-4，因此对 bfloat16 的整像素平移断言本就不可达。这些条目以 known-failure / xfail 方式记录在 testing/known_failure_xfails 与 testing/half_precision_xfails 中，属于测试预期管理而非回归。

六、迁移与工程启示

对本变更的使用者而言，需要关注以下几点：

1. eager 半精度网格数值会变化：这是 breaking change 的核心——修复让 eager 结果向"正确舍入"一侧移动，如果你的模型在半精度下对网格有紧耦合的数值依赖（例如保存了依赖旧网格数值的权重或校准参数），需要重新验证；
2. 编译与非编译不再分叉：torch.compile包裹模型前后的网格行为从此一致，消除了此前"编译后结果变化"的隐性问题，调试编译图时不再需要怀疑网格来源；
3. 半精度的硬件边界仍在：CUDA 下 float16 网格仍允许最多 1 ulp 的平台差异（ATen 与 Triton 除法不同），跨后端严格对比时应使用 bfloat16 或提升到 float32 工作精度；
4. 断言精度要与 dtype 匹配：如test_translation所示，对 bfloat16 数据设置atol=1e-4是不可达的，半精度测试的容忍度应放宽到 dtype 的 ulp 量级。

作为参考，is_compiling()（kornia/core/utils.py）是 Kornia 判断是否处于torch.compile/torch.export捕获期的统一入口，create_meshgrid的 capture 分支（grid.py）正是依赖它与torch.jit.is_tracing()来区分执行模式——理解这一守卫逻辑，有助于把握 Kornia 中大量"eager 与编译路径分叉"代码的设计范式。

• 计算机视觉
• 深度学习
• 图像处理