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

本篇技术指南围绕 Kornia 中nms2d（2D 非极大值抑制）修复展开：该函数此前仅能接受方形kernel_size（如(3, 3)、(5, 5)），一旦传入(11, 5)这类非方形窗口就会抛出RuntimeError: shape '[1, 1, -1, H, W]' is invalid for input of size ...。修复后，nms2d可正确接受任意非方形kernel_size。读完本文，你将理解该错误的根因（邻域核两维度被交换、padding 应用到错误的边）、修复后的通用路径实现原理、边界语义约定，以及如何用测试用例验证修复的正确性。

背景：nms2d 在 Kornia 中的角色

nms2d是 Kornia 空间几何子包 kornia/geometry/subpix/ 中的核心函数，用于把稠密的响应图（response map）转化为稀疏的局部极大值候选点。它的典型调用方式是把角点/关键点的响应图输入进去，仅保留那些在kernel_size窗口内严格大于所有邻居的位置，其余位置置零（或返回布尔掩码）。

在 Kornia 的局部特征检测流水线中，nms2d被广泛使用，例如：

• kornia/feature/disk/detector.py 中对 heatmap 做 NMS 提取关键点；
• kornia/feature/aliked/aliked.py 中对检测分数做 NMS；
• kornia/feature/xfeat.py 中nms2d(x, (kernel_size, kernel_size), mask_only=True) & (x > threshold)组合掩码与阈值筛选。

这些内部调用大多使用方形窗口（(window_size, window_size)），但 API 层面kernel_size的类型签名是tuple[int, int]，并没有强制要求两个维度相等——这正是 bug 的伏笔。

问题现象：非方形窗口触发view形状错误

修复前，向nms2d传入非方形kernel_size（例如(11, 5)）会直接抛错：

RuntimeError: shape '[1, 1, -1, H, W]' is invalid for input of size ...

这个错误信息中的[1, 1, -1, H, W]揭示了问题本质：内部实现把邻域窗口的卷积核 reshape 时，两个 spatial 维度被交换（即H、W位置上出现了-1这种无法推断的维度），同时 padding 被应用到了错误的边对（edge pair）上，导致卷积输出的空间尺寸与后续view期望的尺寸对不上。

该问题由 #4240 记录并提出，#4242 完成修复。这是 Kornia 变更日志体系（changelog.d/+migration-082.fixed.md）中记录的众多修复之一，同类问题还出现在 3D 版本（[nms3d曾对除(3, 3, 3)外的所有核尺寸抛错，见 #4241，相关测试注释记录于 tests/geometry/subpix/test_nms.py）。

根因分析：卷积核两维交换与 padding 错边

旧实现基于卷积（convolution）+ padding 构造邻域比较。它先创建一个形状为(k_y * k_x, 1, k_y, k_x)的核，再对输入做 padding 后卷积。代码层面存在两个叠加的缺陷：

1. 两个 extents 被交换：核的构建过程中把k_y和k_x弄反，导致卷积核的宽高与预期相反；
2. padding 应用到了错误的边：padding 本应按(k_y - 1) // 2（上下）和(k_x - 1) // 2（左右）分别施加，却被打到了错误的维度组合上。

当k_y == k_x（方形窗口）时，这两个缺陷恰好相互掩盖、不产生可见影响；一旦传入非方形窗口，交换后的核尺寸与 padding 后的输入无法匹配，卷积输出形状错误，随后的view便抛出文章开头那条shape '[1, 1, -1, H, W]' is invalid的RuntimeError。

这条根因线索也直接写进了测试注释中：tests/geometry/subpix/test_nms.py 的test_non_square_kernel_size明确说明"非方形窗口过去会因卷积核两维交换、padding 应用到错误边而在内部view处抛错（#4240）"。

修复后的实现：无 padding 的通用路径

修复后的实现位于 kornia/geometry/subpix/nms.py，彻底抛弃了"卷积 + padding"方案，改为分区 max-pooling 累加的通用路径（general path），从而天然支持任意(k_y, k_x)组合。核心有三块：

1. 邻域范围拆分：_split_window

kornia/geometry/subpix/nms.py#L27-L30 的_split_window把长度为k的窗口沿中心拆成"中心前（before）"和"中心后（after）"两个范围：

def _split_window(k: int) -> Tuple[int, int]: """Return the neighbour extents of a length-``k`` window either side of its centre.""" before = (k - 1) // 2 return before, k - before - 1

例如k = 5返回(2, 2)，k = 11返回(5, 5)。它统一了奇数/偶数窗口的中心定义，并且不依赖卷积核形状，因此对(11, 5)、(5, 11)、(13, 9)等任意组合都成立。

2. 邻域最大值计算：_neighbourhood_max2d

kornia/geometry/subpix/nms.py#L41-L68 的_neighbourhood_max2d计算"每个k_y x k_x窗口内、排除中心后的最大值"。它把窗口去掉中心后划分为四个矩形区域（中心行上方、中心行下方、中心行左/右侧），分别用F.max_pool2d归约，再通过torch.maximum逐对合并（见_reduce_max）。由于两个全宽条带共享一次(1, kx)的列向池化，每个位置的代价是O(k_y + k_x)次访问，而不是朴素的k_y * k_x - 1次——即便支持任意核尺寸，性能依然可控。

3. 边界语义：窗口放不下的位置不是极大值

通用路径 kornia/geometry/subpix/nms.py#L277-L288 中，没有任何 padding：

else: # General path: the same rule as the explicit paths above, for any window size. A # position whose full window does not fit inside the image is not a maximum -- it has # not been compared against the neighbours that would decide it -- so the border strip # stays False and only the interior is computed (#4239). ky, kx = self.kernel_size cy, by = _split_window(ky) cx, bx = _split_window(kx) mask = torch.zeros(B, CH, H, W, device=x.device, dtype=torch.bool) if H >= ky and W >= kx: centre = x[..., cy : H - by, cx : W - bx] mask[..., cy : H - by, cx : W - bx] = centre > _neighbourhood_max2d(x, ky, kx)

这里确立了一条统一的边界约定（#4239）：位于距边界(k - 1) // 2以内的位置永远不是极大值，因为它的完整窗口放不进图像、无法与所有邻居比较。(3, 3)、(5, 5)、(7, 7)三条手写快速路径一直遵循此约定，修复后通用路径也与之保持一致。

函数入口与模块/函数式 API 的一致性

nms2d函数式入口定义于 kornia/geometry/subpix/nms.py#L390-L395，本质是NonMaximaSuppression2d模块的一行包装：

def nms2d(input: torch.Tensor, kernel_size: tuple[int, int], mask_only: bool = False) -> torch.Tensor: r"""Apply non maxima suppression to filter. See :class:`~kornia.geometry.subpix.NonMaximaSuppression2d` for details. """ return NonMaximaSuppression2d(kernel_size)(input, mask_only)

NonMaximaSuppression2d模块（nms.py#L103-L292）支持两个参数：

• kernel_size: tuple[int, int]：邻域窗口大小，构造时校验必须是长度为 2 的 tuple；
• forward(x, mask_only=False)：输入形状(B, C, H, W)；mask_only=True时返回布尔极大值掩码，否则返回非极大值置零后的响应图。

此外，该模块重写了_load_from_state_dict（nms.py#L117-L133），把旧卷积实现遗留的kernel持久化 buffer 从 state_dict 中弹出，保证旧版本保存的 checkpoint 在strict=True下仍能加载——这一点由测试test_loads_legacy_kernel_state_dict_strictly（tests/geometry/subpix/test_nms.py）覆盖。

测试验证：参考实现与参数化回归

修复正确性由一套"从定义出发、与实现零共享结构"的参考实现兜底。tests/geometry/subpix/test_nms.py#L28-L54 的_reference_nms_mask通过平移整块视图、torch.maximum归约，逐定义写出"严格大于窗口内其他所有值"的规则，再与nms2d输出逐元素对比。

针对本次修复的关键回归测试是test_non_square_kernel_size（tests/geometry/subpix/test_nms.py#L233-L243）：

@pytest.mark.parametrize("kernel_size", [(11, 5), (5, 11), (13, 9)]) def test_non_square_kernel_size(self, device, dtype, kernel_size): # A non-square window used to raise out of an internal `view`, because the conv kernel was # built with its two extents swapped and the padding was applied to the wrong pair of # edges (#4240). torch.manual_seed(0) inp = (torch.rand(1, 2, 29, 31, device=device) * 4).to(dtype) self.assert_close( kornia.geometry.subpix.nms2d(inp, kernel_size, mask_only=True).to(dtype), _reference_nms_mask(inp, kernel_size).to(dtype), )

它同时覆盖了"高度大于宽度"（(11, 5)）、"宽度大于高度"（(5, 11)）和"非对称奇数"（(13, 9)）三种情况，并与参考实现做全量比对。配套的还有一组行为契约测试，共同锁定 NMS 语义：

实战用法

1. 非方形 NMS：保留横向/纵向不同的抑制半径

在特征检测中，如果希望沿水平方向抑制得更宽、沿垂直方向抑制得更窄（例如检测拉长的结构），现在可以直接传入非方形窗口：

import torch from kornia.geometry.subpix import nms2d, NonMaximaSuppression2d # 响应图: (B, C, H, W) scores = torch.rand(1, 1, 32, 32) # 函数式 API：水平半径 5、垂直半径 2 的非方形窗口 mask = nms2d(scores, (11, 5), mask_only=True) print(mask.dtype, mask.shape) # torch.bool, torch.Size([1, 1, 32, 32]) # 模块式 API：同样的窗口，返回抑制后的响应图 nms = NonMaximaSuppression2d((11, 5)) suppressed = nms(scores) # 非极大值位置被置零

2. 与阈值筛选组合（XFeat 风格）

参考 kornia/feature/xfeat.py#L346 的写法，把 NMS 掩码与置信度阈值取交集，得到最终稀疏关键点位置：

mask = nms2d(scores, (7, 7), mask_only=True) & (scores > 0.5) ys, xs = torch.nonzero(mask, as_tuple=False).unbind(1) # 注意去掉 batch/channel 维度

3. 模型加载兼容性

若你保存过旧版本 Kornia 训练的包含NonMaximaSuppression2d的模型（如嵌入在ScaleSpaceDetector中，见 kornia/feature/scale_space_detector.py#L765），升级后无需重新保存：_load_from_state_dict会静默接受旧的kernelbuffer，strict=True加载不会报unexpected key。

小结

本次nms2d修复（changelog.d/+migration-082.fixed.md，#4240/#4242）从根上消除了"卷积核两维交换 + padding 错边"这一对仅在非方形窗口下暴露的缺陷，用_split_window+_neighbourhood_max2d的无 padding 通用路径取代了卷积实现，同时把边界条带永不判极值的语义在快速路径与通用路径之间统一（#4239）。配套的参数化回归测试（(11, 5)、(5, 11)、(13, 9)）保证了修复与参考定义严格一致。如今nms2d的kernel_size参数名副其实地支持任意二维窗口，为检测各向异性结构的关键点提供了更自由的抑制半径选择。相关实现与测试可直接在 kornia/geometry/subpix/nms.py 与 tests/geometry/subpix/test_nms.py 中查阅，API 文档见 docs/source/geometry.subpix.rst。

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