Kornia LAF 补丁提取的 CPU 半精度采样修复：float32 回退与批处理 grid

计算机视觉
深度学习
人工智能
图像处理

【免费下载链接】kornia

🐍 空间人工智能的几何计算机视觉库

导读

本文围绕 changelog.d/+migration-110.fixed.md 记录的一次关键修复展开：在 PyTorch ≤ 2.9 中，float16/bfloat16的 CPUgrid_sample内核在旋转补丁跨越图像边界时存在越界读取缺陷，会返回 NaN 或远超图像取值范围的垃圾值。Kornia 的extract_patches_simple与extract_patches_from_pyramid通过「在所有设备上以 float32 采样再转回半精度」的方式绕开该缺陷，同时把逐图像 Python 循环折叠为一次批处理的grid_sample，并在torch.compile(fullgraph=True)下消除了随 batch 增长的计算图膨胀。读完本文，你将掌握这两个 LAF 补丁提取器的精度策略、内存分块机制以及它们与 torch.compile 的交互原理。

背景：两个 LAF 补丁提取函数

Kornia 的局部特征（LAF，Local Affine Frames）管线中，extract_patches_simple与extract_patches_from_pyramid是核心的采样算子，均定义于 kornia/feature/laf.py（导出入口见 kornia/feature/init.py）：

extract_patches_simple(img, laf, PS=32, normalize_lafs_before_extraction=True)（实现）：直接按 LAF 定义的仿射框从原图采样，不做平滑，因此有较强混叠（aliasing），适合快速提取。
extract_patches_from_pyramid（实现）：从图像金字塔的合适层级采样，尺度小于PS的 LAF 会被路由到仍能提供完整补丁的最粗层级，避免混叠。

两个函数的输入输出约定一致：图像为(B, CH, H, W)，LAF 为(B, N, 2, 3)，输出补丁为(B, N, CH, PS, PS)。它们都是纯函数式 API，供LAFOrienter、LAFAffNetShapeEstimator、LAFDescriptor等下游模块复用（相关修复链条见 changelog.d/+migration-111.fixed.md）。

问题根因：torch ≤ 2.9 的半精度 CPU 越界读取

本次修复的核心是一个上游（PyTorch）缺陷：当旋转后的补丁跨越图像边界时，torch ≤ 2.9 中float16/bfloat16的 CPUgrid_sample内核会越界读取（out-of-bounds read），从而返回 NaN 或远超出图像取值范围的垃圾值（例如数量级1e4的随机堆数据）。原因在于：

旋转补丁的采样网格有相当一部分落在图像外，这些坐标由边界填充（padding_mode="border"）来补值；
半精度 CPU 内核在这些边界坐标上没有正确处理，读取了非法内存。

该问题在测试中专门以「补丁必须保持在图像取值范围内」这一不变式来刻画。见 tests/feature/test_laf.py 的test_border_patches_stay_in_range：任何跨界补丁采到的值都应是图像像素值的凸组合，因此必须落在[img.min(), img.max()]之内且有限（isfinite）。而缺陷内核会返回零、1e4量级数值或 NaN，取决于堆内存中恰好残留的内容。

修复方案一：全设备 float32 采样回退

修复的核心策略是把半精度输入在采样前统一提升到 float32，采样完成后再把补丁转回原始精度。从源码看，这一逻辑体现在两条关键链路上：

网格精度提升：_grid_dtype将float16/bfloat16一律映射为torch.float32；_promoted_grid_dtype进一步对图像与 LAF 做torch.promote_types后应用该策略，保证坐标运算不丢失任何一侧的精度。
图像一次上采样：在 extract_patches_simple 的实现中，sample_img = img.to(grid_dtype)在分块循环之前完成，避免每个块重复对整幅图做类型转换；_grid_sample_patches（实现）保证img与grid共享 dtype，直接调用F.grid_sample。

值得注意的是，该上采样并非只针对 torch ≤ 2.9 的坏内核——_grid_sample_patches的 docstring 明确指出这是在所有 torch 版本上有意为之：半精度采样坐标在大图像上会量化为整像素（即归一化坐标精度损失），float32 网格计算可避免这一损失。_grid_dtype的 docstring 也说明了它对 float16/bfloat16 统一提升的原则。

从代码注释看，extract_patches_from_pyramid同样遵循该约定：金字塔 atlas 直接用网格 dtype 构建（kornia/feature/laf.py#L795-L807），使 replicate 填充、pyrdown与每个分块的grid_sample都运行在所有 torch 版本都稳定的内核上。

精度与性能的权衡

源码明确记录了在 CUDA 上的代价：原生半精度核在 CUDA 上是正常的（不越界），但为了统一的正确性策略，float16 的extract_patches_simple在高 N 场景下大约付出2 倍减速，换取准确性。这在 migration-110 的 breaking changes 说明中被称为 "CUDA cost"——即该修复改变了所有后端的半精度补丁数值，而不仅是 CPU。

混合精度 autocast 管线的兼容

测试test_mixed_dtype_laf_under_autocast（tests/feature/test_laf.py#L753-L766）验证了典型的 autocast 场景：检测器（如 KeyNet）在 autocast 下输出 half/bfloat16 的 LAF，而源图仍是 float32。提取器自己完成小张量（LAF）的类型提升，而不是拒绝这一合法管线或依赖后端各自的grid_sample隐式提升。对应的test_mixed_dtype_preserves_laf_precision（tests/feature/test_laf.py#L768-L778）则反向验证：float32 LAF 配 half 图像时，LAF 的亚像素坐标不能被降成 half 再提升回来——这正是网格使用提升后 dtype 的原因。

修复方案二：折叠的批处理 grid_sample

此前两个提取器采用「逐图像 Python 循环」，每张图像调用一次grid_sample。本次修复将逐图像循环替换为折叠的批处理grid_sample：

每个 LAF 的(PS, PS, 2)网格被折叠为(B, N*PS, PS, 2)，一次调用覆盖整个 batch。该逻辑见_sample_patches：folded = grid.view(B, N * PS, PS, 2)后单次_grid_sample_patches，结果再 reshape 为(B, ch, N, PS, PS)并 permute 成(B, N, ch, PS, PS)。
当N很大时，网格与采样结果按N方向分块（chunking），以约束工作区（workspace）峰值内存。分块大小由_grid_chunk_lafs决定：它按B * PS * PS * max(2, ch) * elem_size估算每个 LAF 的网格与通道缩放采样结果的字节数，用默认 64 MiB 预算求最大可容纳的 LAF 数；当一切都装得下时，循环退化为单次调用的快路径。

测试对该语义做了三重验证：

test_chunked_matches_single_call（tests/feature/test_laf.py#L857-L870）通过 monkeypatch 把_grid_chunk_lafs强制为 1（每块一个 LAF），断言与单次调用结果完全一致（CPU 上逐位相等）；
test_chunk_budget_accounts_for_channels（tests/feature/test_laf.py#L872-L878）验证高通道特征图会把块数压小：ch=1时 1000 个 LAF 一次装下，ch=256时每块仅 64 个；
test_channel_laf_correspondence（tests/feature/test_laf.py#L880-L892）逐 LAF 对比单块提取结果，锁定ch>1且N>1时通道与 LAF 的对应关系。

对 torch.compile(fullgraph=True) 的图优化影响

这是本次修复中容易被忽略却影响深远的收益。迁移说明中给出了具体的实证：

旧实现（逐图像循环）在 trace 时会把循环展开成每个 batch 元素一个grid_sample，计算图随 batch 增长，且每个新 batch 尺寸都会触发一次重编译。例如分别用 batch 2、3、5、7 追踪extract_patches_simple，会构建出含 2/3/5/7 个grid_sample节点的 4 张图；新实现只需构建各含 1 个节点的 2 张图。

换句话说：

旧图：图的规模 = O(batch)，batch 变化即重编译；
新图：批处理折叠后图固定为 1 个grid_sample节点，图的规模与 batch 解耦。

两种形式都能在fullgraph=True下编译，但只有新实现避免了「图随 batch 线性膨胀 + 每次新 batch 重编译」的开销。测试test_dynamo（tests/feature/test_laf.py#L926-L935）直接以torch.compile(..., fullgraph=True)断言提取器可以完整 trace 为单张图。

无数据依赖分支的实现细节

为了保证fullgraph路径可编译，两个提取器在非有限 LAF 处理上也刻意保持「无条件」：extract_patches_simple在循环结束后统一masked_fill_清零（kornia/feature/laf.py#L713-L715），extract_patches_from_pyramid则对pyr_idx < 0的条目做同样的无条件清零（kornia/feature/laf.py#L856-L872），避免数据相关的 Python 分支导致图断裂。

附带修复：非有限 LAF 与 MPS 边界行为

本次 migration 涉及的同族修复还包括：

migration-111：任何含 NaN/Inf 的 LAF 帧（哪怕只有中心一个元素）在网格算术前被整体标记并净化，返回全零补丁与零 LAF 梯度，而不是把非法网格交给grid_sample——其 CPUgrid_sampler_2d_backward边界填充内核可能直接终止进程。测试见 tests/feature/test_laf.py#L780-L810。
migration-126：MPS 没有padding_mode="border"，用「零填充 + 网格截断」模拟时，截断目标从±1（align_corners=False下是边界像素的外边缘，会与零填充混出约一半的暗色值）修正为最外层像素中心±(1 - 1/size)，使跨界补丁与 CPU 的最大偏差从0.395降到2.5e-6。

实际使用建议

无需用户侧改动：float32回退对调用方透明，输入 half 图像仍返回 half 补丁（dtype 不变，见test_border_patches_stay_in_range中对patches.dtype == half_dtype的断言）。
设备与 dtype 约定：输出补丁始终落在图像所在设备与 dtype 上；LAF 允许与图像不同设备/精度，提取器会先迁移（kornia/feature/laf.py#L684-L688）。测试test_laf_on_another_device（tests/feature/test_laf.py#L914-L924）验证了该契约。
大 batch 与大图：分块机制默认以 64 MiB 预算约束峰值内存；对超高通道特征图（如ch=256）建议预判块数（可用_grid_chunk_lafs估算），避免意外的工作区压力。
训练管线：若你的检测器在训练中可能产出退化的非有限 LAF，本次修复已保证此类帧得到零补丁与安全反向传播，无需额外防御逻辑。
编译部署：若使用torch.compile(fullgraph=True)，批处理折叠后的提取器图规模与 batch 解耦，重编译次数大幅下降，可放心放入编译后的推理/训练图。

总结

extract_patches_simple与extract_patches_from_pyramid的这次修复同时解决了三个层面的问题：正确性（绕开 torch ≤ 2.9 半精度 CPU 越界读取）、精度（避免大图像上归一化坐标的精度损失）、性能/可编译性（批处理折叠 + 分块内存约束 + 消除随 batch 膨胀的计算图）。其实现细节——提升后 dtype 的一次性上采样、无数据依赖分支的清零策略、64 MiB/128 MiB 的分块预算——都可在 kornia/feature/laf.py 与其配套测试 tests/feature/test_laf.py 中逐一验证，是理解 Kornia 局部特征管线精度与性能权衡的绝佳案例。