Kornia 修复深度解析：HyNet 与 SOSNet 半精度描述符的 CPU/GPU 稳定性改造

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

本文基于 Kornia 仓库 changelog.d/+migration-085.fixed.md 中的修复记录，深入解析HyNet与SOSNet两个局部图像描述符网络在半精度（float16/bfloat16）下无法在 CPU 运行、且退化 patch 输出全 NaN 的两大缺陷：一条缺陷源自avg_pool3d在 CPU 上没有半精度 kernel，另一条源自1e-10的eps在float16中不可表示而冲刷为 0。读完本文，你将理解该修复的设计决策、实现细节（sosnet.py、hynet.py）、对应的测试保障，以及该变更对半精度描述符数值行为的影响。

修复背景：一条迁移记录里的两个缺陷

在 Kornia 的 changelog 驱动开发流程中，changelog.d/目录下的每条+migration-*.fixed.md记录了一个已合入的缺陷修复。本次记录（对应 issue #4224，PR #4225）明确说明：

HyNet和SOSNet现在可以在半精度下运行——无论是 CPU 还是 GPU——并且对退化 patch 不再返回 NaN。

修复前，两个缺陷"坐在同一行代码上"（two defects sat on the same line），即位于两个网络forward末尾的LocalResponseNorm归一化步骤：

• CPU 半精度 kernel 缺失：两个网络的最终LocalResponseNorm收到的是一个 4 维(B, C, 1, 1)张量，而torch.nn.functional.local_response_norm对 4 维输入会路由到avg_pool3d，该 kernel 在 CPU 上没有float16/bfloat16实现，因此无论输入是什么，CPU 上都会直接抛出：

  NotImplementedError: "avg_pool3d_out_frame" not implemented for 'Half'

  （'BFloat16'拼写同样不存在该实现）

• float16下eps冲刷为 0 导致 0/0：在float16且 kernel 存在的设备（如 MPS、CUDA）上，只要网络把某个 patch 映射到恰好为零的激活，描述符就会全部变成 NaN。

缺陷一：local_response_norm的avg_pool3d路由与 CPU kernel 缺失

LocalResponseNorm在描述符网络中的角色

先看两个网络的归一化步骤在源码中的位置：

• SOSNet.forward：特征提取层输出descr后，self.desc_norm = nn.Sequential(nn.LocalResponseNorm(256, alpha=256.0, beta=0.5, k=0.0))，即对每个描述符向量执行 L2 归一化。
• HyNet.forward：self.desc_norm = nn.LocalResponseNorm(2 * self.dim_desc, 2.0 * self.dim_desc, 0.5, 0.0)，同样在末尾做归一化。

这两个desc_norm的输入形状都是(B, C, 1, 1)——特征图已经过kernel_size=8的全局卷积（SOSNet 的最后一个Conv2d(128, 128, kernel_size=8)、HyNet 的Conv2d(128, dim_desc, kernel_size=8)）压缩为单像素。而 PyTorch 的F.local_response_norm对 4 维输入会通过avg_pool3d计算邻域平方和，这正是缺陷所在：avg_pool3d的 CPU kernel 对半精度数据类型没有注册实现，于是 CPU 上所有半精度调用都会撞上NotImplementedError。

这也是为什么修复记录强调该归一化步骤与kornia.feature.siftdesc中给eps守卫的保护不同：siftdesc的守卫只针对float16做提升（原因见缺陷二），而此处由于 CPU kernel 缺口同时覆盖float16与bfloat16两种 dtype，必须做一个更宽的提升。

缺陷二：1e-10的eps在float16中不可表示

常数 patch 如何一路走到"恰好为零"

eps的作用是保证归一化的除法在描述符向量恰好为零时依然有定义：

• SOSNet.forward(input, eps: float = 1e-10)：sosnet.py 中eps在descr + eps中加入，desc_norm内部用该值做除法分母的守卫。
• HyNet.forward(x)：hynet.py 中eps_l2_norm: float = 1e-10，同样在desc_norm的除法中起守卫作用。

关键事实是：这两个网络的任意常数 patch 都会精确走到零激活：

• SOSNet 中每个Conv2d都是bias=False，每个BatchNorm2d都是affine=False（见 sosnet.py 的nn.Sequential结构），因此一个常数 patch 经过零偏置卷积与无仿射参数的 BatchNorm 后，到达归一化层的值就是精确的 0。
• HyNet 只有在is_bias=False配置下才会走到零（hynet.py）：因为 HyNet 使用FilterResponseNorm2d+TLU（阈值线性单元）组合，TLU 的偏置开启时，归一化前的张量永远不会精确落在 0 上。

float16与bfloat16的本质差异

修复记录给出了精确的数值解释：

• float16的指数范围有限，1e-10不可表示，会被冲刷（flush）为0.0，于是除法变成0/0，产生 NaN。这是 MPS/CUDA 等 kernel 存在的设备上出现全 NaN 的根因。
• bfloat16保留了float32的指数范围（只是尾数精度更低），因此1e-10在bfloat16中依然可表示，守卫在这里是存活的，bfloat16只受 CPU kernel 缺失问题影响。

这也解释了为何siftdesc的_rootsift（siftdesc.py）只对float16做提升：SIFT 的守卫面临的是同样的eps不可表示问题，但不存在 CPUavg_pool3d的 dtype 覆盖差异——bfloat16那里没有1e-10冲刷问题，所以只需要保护float16。

修复方案：归一化步骤提升到float32再回写

修复的核心决策是：无论输入是float16还是bfloat16，都只把最后那一步归一化提升到float32计算，再转回原 dtype。两个网络的具体实现：

# SOSNet.forward 的归一化分支（sosnet.py） if descr.dtype in (torch.float16, torch.bfloat16): descr = self.desc_norm(descr.float() + eps).to(descr.dtype) else: descr = self.desc_norm(descr + eps)

# HyNet.forward 的归一化分支（hynet.py） if x.dtype in (torch.float16, torch.bfloat16): x = self.desc_norm(x.float() + self.eps_l2_norm).to(x.dtype) else: x = self.desc_norm(x + self.eps_l2_norm)

设计要点：

1. 提升范围最小化：只提升归一化这一步，而不是整个网络前向。网络的卷积、归一化主体仍以半精度执行，从而保留半精度的性能收益。
2. 覆盖两种 half dtype：float16解决eps冲刷问题，bfloat16解决 CPU kernel 缺失问题，一次dtype in (torch.float16, torch.bfloat16)的分支同时覆盖两者——比siftdesc的float16单独提升更宽，原因正如前文所述：CPU kernel 缺口对两种 dtype 都成立。
3. float32/float64完全不动：这两种精度走原始表达式，在 CPU、CUDA、MPS 上逐位（bitwise）不变，即常规精度用户不会观察到任何行为变化。

注意HyNet的提升分支中x.float() + self.eps_l2_norm的加法发生在float32域内，确保1e-10以float32语义参与运算而非先被截断。

数值行为变化：半精度输出会变，但"变得更准"

修复并非零成本：半精度描述符的输出会改变，因此该变更是一次有意的数值语义调整，而不是纯粹的 bug 修复。changelog 给出了可验证的量化结论：

• 输出变化幅度为0.25–2.25 eps（eps即 dtype 的单位舍入误差量级）。
• 变化最大的方向是朝向float64参考值移动，而不是仅仅远离 NaN：
  • SOSNet在float16下对float64同权重模型的最大绝对误差从2.15e-03降到2.28e-04，即误差下降约一个数量级。
• 未改进的配置也保持在原值一个eps之内，即不会引入明显退化。

由此带来的兼容性提示：半精度描述符在本版本前后不再逐位可比。如果你的应用依赖跨版本的半精度描述符做缓存、数据库索引或位级比对，升级后需要重新生成。

测试保障：退化 patch 与编译路径的双重守护

该修复在测试层有多重验证（test_sosnet.py、test_hynet.py）：

1. 退化 patch 有限性测试（直接回归守护）

# test_sosnet.py: test_degenerate_patch_gives_finite_descriptors @pytest.mark.parametrize("patch_value", [0.0, 0.5]) def test_degenerate_patch_gives_finite_descriptors(self, device, dtype, patch_value): patches = torch.full((2, 1, 32, 32), patch_value, device=device, dtype=dtype) sosnet = SOSNet().to(device, dtype).eval() descriptors = sosnet(patches) assert torch.isfinite(descriptors).all(), ...

对应地，test_hynet.py 的test_degenerate_patch_gives_finite_descriptors额外参数化is_bias为True/False，覆盖 HyNet 走到精确零（is_bias=False）与不走零（is_bias=True）两种路径。两个测试都使用torch.full构造常数 patch，并在全部设备 × dtype 组合下断言输出无 NaN/Inf。

2. 编译路径守护

两个测试类都包含test_dynamo、test_dynamo_fullgraph（torch.compile(..., fullgraph=True)）与test_jit。测试注释明确指出：归一化的 dtype 分支读取的是静态元数据，torch._dynamo在 trace 时即可解析该分支，因此编译后的结果必须与 eager 一致——包括被提升的半精度路径（若分支被错误丢弃会重新返回 NaN）。同时docs/source/_data/export_support.json中记录了feature.HyNet与feature.SOSNet为可捕获（capturable）模型，本次修复不能引入 graph break。

使用与验证：如何在自己环境中复现

复现修复前后差异的最小脚本（基于仓库中公开的 API，kornia/feature/__init__.py中SOSNet、HyNet、FilterResponseNorm2d均已导出）：

import torch from kornia.feature import SOSNet, HyNet for dtype in (torch.float16, torch.bfloat16, torch.float32, torch.float64): patches = torch.full((2, 1, 32, 32), 0.5, dtype=dtype) sosnet = SOSNet().to(dtype).eval() descs = sosnet(patches) print(f"SOSNet {dtype}: finite={torch.isfinite(descs).all().item()} shape={tuple(descs.shape)}") hynet = HyNet(is_bias=False).to(dtype).eval() descs = hynet(patches) print(f"HyNet {dtype}: finite={torch.isfinite(descs).all().item()} shape={tuple(descs.shape)}")

在修复前，float16的SOSNet/HyNet输出会全部为 NaN，CPU 上的两种 half dtype 都会直接抛NotImplementedError；修复后所有组合都应输出finite=True的(B, 128)描述符。更严谨的对照可参考仓库测试：在float64下用同一套随机初始化的模型权重跑一次作为参考（test_gradcheck中即使用torch.float64输入），再比较半精度输出的最大绝对误差。

总结

本次修复（#4224 / #4225）揭示了深度学习中一个典型的"隐性精度陷阱"组合：框架 kernel 覆盖不完整（CPUavg_pool3d无 half 实现）+ 浮点守卫值不可表示（1e-10在float16中冲刷为 0），而两个问题恰好落在同一行归一化代码上。修复以"仅提升最后一步归一化到float32"的最小侵入方案同时解决两者，并通过常数 patch 有限性测试、dynamo/JIT 编译一致性测试建立长期守护。对于使用半精度推理加速的视觉应用，这意味着HyNet/SOSNet描述符管线现在可以安全地在 CPU 与 GPU 上以float16/bfloat16运行，代价是半精度输出相对旧版本不再逐位可比——迁移时建议按本文的量化说明重新评估数值一致性需求。

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