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

中值模糊是图像去噪中的经典非线性算子，但在深度学习框架中，它的实现方式直接决定了小窗口（3x3/5x5）下的吞吐能力。本文基于 Kornia 仓库中的 median_parallelism.md 并行化分析笔记，结合 kornia/filters/median.py 源码与 tests/filters/test_median.py 测试，完整讲解 Kornia 中值模糊的标准路径、CPU 固定选择网络快速路径的底层原理、语义约束（平局索引、NaN/Inf、梯度）、PyTorch 上游实现现状，以及"移动轴、改变内存布局"实测结论与 channels-last 的实战取舍。读完本文，你将理解小窗口median归约为何是性能瓶颈、Kornia 为何保留并仅限推理期启用 CPU 快速路径、以及如何在自己的图像 pipeline 中通过内存布局优化获得收益。

背景：中值模糊的标准路径与并行化窗口

Kornia 中值模糊的标准（通用）路径位于 kornia/filters/median.py：先用二值卷积核把每个局部窗口展开成特征轴，再沿该轴取中值。

padding = _compute_zero_padding(kernel_size) kernel: torch.Tensor = get_binary_kernel2d(kernel_size, ...) b, c, h, w = input.shape # 将局部窗口映射为单个向量 features = F.conv2d(input.reshape(b * c, 1, h, w), kernel, padding=padding, stride=1) features = features.view(b, c, ky * kx, h, w) # B x C x (K_h * K_w) x H x W # 沿特征轴取中值 return features.median(dim=2)[0]

其核心形态是：先构造一个B x C x K x H x W的窗口张量，再调用median(dim=2)。对于 3x3 或 5x5 窗口，K=9或K=25，即每个输出像素对应一次对 9 或 25 个值的归约，共有B*C*H*W个相互独立的归约。

从并行化角度看，中值模糊应该沿着输出像素、通道和 batch 元素并行展开——这是天然的"海量独立任务"形态；而把 CPU 线程分配到单个9/25 值窗口内部则过于细粒度，收益极低。GPU 线程可以在窗口内协作，但这种映射仍然需要最小化每个像素的工作量与内存流量。

当前 PyTorch 的median(dim)实现：三个后端的观察

分析笔记针对 PyTorch 提交fc553088记录了 CPU、CUDA、MPS 三个后端对median(dim)的底层实现观察。下表是原文核心结论的完整整理，注意这些都是实现层面的观察，只有明确标注的才是上游实测数据，不能替代基准套件的结论。

Kornia 的 CPU 快速路径：固定选择网络

针对上述"通用选择算法对小窗口不成比例"的问题，Kornia 在 kornia/filters/median.py 中实现了固定选择网络（fixed selection network）：

• _median_network(size)将 Batcher 奇偶归并排序网络（Batcher, "Sorting networks and their applications", AFIPS 1968）剪枝到中间输出线，缺失的线用符号正无穷表示。剪枝后只保留影响中值输出位置的比较器对。
• _MEDIAN_NETWORKS = {3: _median_network(9), 5: _median_network(25)}，即只为 3x3 与 5x5 窗口预生成网络。
• _median_blur_network(input, size)不物化 patch、不排序：先 pad，把窗口内每个偏移的视图放入values列表，再按预生成的比较器对执行torch.minimum/torch.maximum逐层比较，最后取values[size * size // 2]。

网络路径还保留了 NaN/Inf 语义：由于原始 one-hot 卷积会把窗口内任意 NaN/Inf 传播到每个提取的特征（包括0 * Inf），快速路径改为对有限值掩码做最大池化得到invalid掩码，再用torch.where(invalid, nan, selected)恢复与标准路径一致的 NaN 传播行为（见 median.py）。

median_blur在 median.py 中的分派条件完整如下：

• 设备为 CPU，或 CUDA 且is_compiling()或ky == 3且输入元素数>= 1 << 20（大输入下才值得用 eager 网络）；
• 未开启 autocast；
• 通道数非 0；
• 输入为浮点、不要求梯度（not input.requires_grad）；
• 无 forward-mode AD tangent（unpack_dual(input).tangent is None）；
• 核高宽相等且ky在_MEDIAN_NETWORKS中（即 3 或 5）。

不满足条件时回退到标准卷积 +median(dim=2)路径；导出（ONNX）场景则用features.sort(dim=2)取同一（lower-middle）元素，因为median.dim没有 ONNX lowering（median.py）。

语义约束：为什么快速路径只能用于推理

median(dim)同时返回值和索引，索引决定反向传播路径。这是限制快速路径最根本的语义因素：

• CPU 的当前比较器对相等值按原始索引排序；CUDA 文档则把选中的平局索引声明为不确定（对应 PyTorch issue #161841）。
• min/max 选择网络返回相同的中值值，但无需保持该索引及其梯度。
• 因此快速路径必须严格限定在推理场景：现有 Kornia CPU 快速路径排除了 autograd 和 forward-mode AD（分派条件中的requires_grad与tangent检查）。
• 同时它必须保持 PyTorch 的lower-middle 选择（偶数窗口取偏下中位）与NaN/Inf 传播语义。

这些约束在 tests/filters/test_median.py 中被逐一验证：test_selection_matches_convolution对照卷积+median 参考路径（覆盖 contiguous / transposed / channels_last 三种布局），test_selection_nonfinite与test_nonfinite_autograd_fallback检查 NaN/Inf，test_tied_gradients_unchanged与test_tied_forward_gradients_unchanged在atol=0, rtol=0下验证普通梯度与 JVP 完全不变，test_selection_dispatch则用 monkeypatch 统计_median_blur_network的调用次数来固定分派规则（CPU 总是选择；eager CUDA 仅 3x3 且元素数达标）。

PyTorch 上游状态

截至文档核对日期（2026-09-18），上游状态如下：

• CPU：torch.median(dim)慢于 sort 的 issue #51450 仍处于打开状态。其报告是历史性能数据，并非对当前 Kornia 分支的实测。
• CUDA：在核对日期的 GitHub 搜索中未发现相关的 PyTorch CUDAmedian性能PR。CUDA 已有专用的通用 median kernel，这是"有实现"的证据，但不代表它对中值模糊已是最优。
• MPS：PR #187060 已于 2026-06-12 合并，把 median/nanmedian 迁移到 Metal 并报告了大输入提速（包括median(dim)场景）；但其结果未覆盖 3x3/5x5 图像窗口，该路径对浮点仍是"排序后收集"。

由此得出的实际方向是：保留现有 CPU 3x3/5x5 推理快速路径，只有在做设备端融合 kernel、并配套聚焦的设备基准与显式的 value/NaN/Inf/gradient 测试时，才继续推进 CUDA/MPS 工作。torch.compile被有意排除在本 eager 路径评估之外——注意源码注释指出 Inductor 会把选择网络融合成单个 CUDA kernel，而 eager CUDA 每个比较器启动一个 kernel，仅在大输入的 3x3 场景下才划算（median.py）。

实测：移动轴与内存布局的影响

文档区分了两个容易混淆的轴：C是图像通道轴，K是包含 9 或 25 个邻域样本的归约轴。

• 用movedim把K移到最后只是改变视图，并不会让K的值在内存中相邻；
• 追加contiguous()才会真正拷贝展开后的窗口张量来做到相邻——而 CPU 的 median 实现本来就会做这种重排。

实测环境：Apple M1、PyTorch 2.14.0、float32、4 个 CPU 线程，输入1 x 3 x 256 x 256、5x5 窗口。数值为三次重复测量的中位数，每次使用common.time_us(min_run_time=0.35)；MPS 在计时区域内同步。每个变体都包含其提取或布局转换开销，输出均与原始卷积/median 路径比对过。

前三行是通用回退路径的诊断变体，不是替代性的公开实现；它们并没有超过当前 CPU 3x3/5x5 推理网络。在额外的 3x3/7x7 与 batch 探针中，移动K并非一致的改进，显式拷贝还在 batched MPS 场景下发生了回退。

物理 channels-last 图像存储可以在不改变 BCHW 形状的前提下帮助现有的 CPU 选择网络：

image = image.contiguous(memory_format=torch.channels_last) result = kornia.filters.median_blur(image, 5)

更广的 CPU 探针覆盖了 batch 1/8、通道 1/3/16/32、图像尺寸 64/128/256、1/4 线程：多通道大图有收益，但转换在 4 线程下回退了 RGB 64x64 场景（5x5 时吞吐约为 0.65x），上面的 MPS 例子也同样回退。因此实现不会强制某种内存格式——在把图像 pipeline 固定为 channels-last 之前，请先用自己的真实形状和线程数做基准。该布局研究没有可用的 CUDA 设备。

如何复现与扩展这类测量

文档强调这些是实现观察，而非对当前分支的权威基准。Kornia 的基准设施位于 benchmarks/：

• benchmarks/common.py 提供统一的测量方法论：warm_up_cpu()让 PyTorch 线程池进入持续负载（混合架构 CPU 上轻载线程会被调度到能效核，未预热会反转 A/B 结论）；time_us()包装torch.utils.benchmark.Timer.blocked_autorange，返回中位数（median）与四分位距（IQR）作为离散度；run_metadata()记录 git commit、平台、Python/torch/kornia 版本、线程数、设备名等可引用元数据；save_json()输出{"metadata": ..., "results": [...]}的严格 JSON。
• benchmarks/filters/flagship.py 是 filters 旗舰基准，median_blur行对比 OpenCVcv2.medianBlur、albumentations、scikit-image、kornia-rs 与 PIL（注意各库边界处理不同：Kornia/scikit-image 用 zero padding，OpenCV 用复制边界）。典型运行方式：

python benchmarks/filters/flagship.py --batches 1,8,32 --size 256 --device cpu python benchmarks/filters/flagship.py --device cuda --compile --json filters_cuda.json python benchmarks/filters/flagship.py --ops gaussian_blur2d,unsharp_mask --compile

结论与实战要点

1. 标准路径：BxCxKxHxW窗口张量 +median(dim=2)，外部并行度充足，但通用选择算法对 9/25 值窗口不成比例。
2. CPU 快速路径：基于 Batcher 网络剪枝出的固定选择网络（median.py），避免物化窗口与排序；受平局索引、梯度与 NaN/Inf 语义约束，仅限推理，且仅对 3x3/5x5 启用。
3. 移动轴无效：movedim/contiguous()调整K轴没有一致收益，显式拷贝在 batched MPS 下反而回退。
4. channels-last 有条件有效：多通道大图可获益（示例 25.61→12.84 ms），但小图/4 线程与 MPS 会回退，务必按实际形状基准后决定；实现不强制任何内存格式。
5. 上游：CPU issue #51450 仍开放，MPS PR #187060 已合并但不覆盖小窗口；进一步的 CUDA/MPS 优化应以设备端融合 kernel + 聚焦设备基准 + 显式 value/NaN/Inf/gradient 测试为前提，torch.compile不在本 eager 路径评估范围内。

对普通 Kornia 用户而言，最简单有效的经验法则是：在 CPU 推理管线中，若图像是多通道且尺寸较大，可以尝试image.contiguous(memory_format=torch.channels_last)后用kornia.filters.median_blur，并以自己的数据跑一次基准再决定是否保留该布局。

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