TorchMorph: CUDA-accelerated Morphological Transforms
Kai Zhao
cs.CV
2026-08-25
上海大学把 22 个形态学、距离变换和 Sinkhorn 算子做成融合 CUDA 核，API 对齐 scipy.ndimage。RTX 4090 上灰度批处理最高快单核 SciPy 约 1100 倍，精确欧氏距离变换约 350 倍。
形态学算子在分割后处理、骨架提取、边界损失里天天用。Python 里的事实标准是 scipy.ndimage：边界模式、结构元素原点、连通性约定，下游库都按它来。它只吃主机内存里的单张 NumPy 数组。
训练循环里的图已经在 GPU 上，而且是一整批。每调一次 ndimage，就要把张量拷回主机、单线程算完再拷回去。现成的 GPU 视觉库覆盖很窄。Kornia 有可微的二维形态学。N 维算子、完整边界模式矩阵和精确欧氏距离变换都不在它的覆盖里。cuCIM 走 CuPy，不是原生 torch.Tensor。MONAI 的若干后处理直接回退到 SciPy 或 CuPy。想同时要批处理 GPU 形态学、精确 EDT、可微的传输距离，得拼三套库、三种张量约定。
上海大学赵凯的 TorchMorph 把这件事收进一个轻量 PyTorch 扩展。作者自己把贡献写成可用性，不是新算法。
对外 22 个算子，分四族：二值形态学、灰度形态学、距离变换、熵正则最优传输。真正手写 CUDA 核的只有腐蚀、膨胀、精确 EDT、chamfer、暴力距离和 Sinkhorn 这几个入口；开运算子、顶帽、梯度、孔洞填充是主机侧用那两个原语拼出来的。新加一种顶帽不用写设备代码。
API 按参数对齐 scipy.ndimage：size / footprint / structure、mode、cval、origin、预分配输出都在。iterations < 1 表示迭代到结果不再变，二值传播和填洞靠这个语义。输入形状是 (B, C, Spatial...)，空间维 1 到 8。移植路径是换 import。
实现分三层。Python 层负责参数归一化和 SciPy 语义，是唯一知道参数约定的地方，结构元素按 structure > footprint > size 解析。pybind11 绑定层只暴露 8 个核入口、6 个 CUDA 翻译单元。核层两条贯穿选择：能在主机上算清的几何就在主机上算清；核按运行时空间维写、编译期封顶，坐标临时量留在寄存器里。
融合形态学核不再对每个输出点、每个结构元素位置做一遍 N 维坐标映射加逐轴边界检查。3 维 3³ 结构元素上，那是每体素 27 次映射。主机把结构元素压成活跃偏移列表，连同按输入 stride 预计算的平坦偏移一起下发，空位不到设备。线程先判断自己是不是内部点：内部点直接把平坦偏移加到线性下标，零边界测试。只有离边界一个结构元素半径之内的线程走通用路径。灰度核覆盖 constant / reflect / nearest / mirror / wrap 五种边界，二值核只用一个 bordervalue。二值核还有提前退出：腐蚀碰到 false、膨胀碰到 true 就可以停，稀疏 mask 上很划算。
精确欧氏距离变换用可分离下包络算法（Felzenszwalb–Huttenlocher）：一维扫描线上对一族抛物线求下包络。扫描本身是串行的，所以一个线程块管一条扫描线，单线程在共享内存里建包络，整块协作加载和查询。并行度来自扫描线数量：B 个 batch、C 个通道、边长 n 的 d 维体，每趟是 B·C·n^{d-1} 条。二维边长不超过 2048 有特化路径，最后的开方融进列扫描；通用路径把当前轴转成最内维保证连续访存；扫描线超出共享内存预算就溢到全局缓冲。chamfer 做可分离前后向扫描，棋盘距离再加对角趟。暴力距离把背景坐标按 256 的瓦片送进共享内存，用来当精确性预言机，也覆盖 chamfer 入口吃不到的各向异性采样。
Sinkhorn 面对的是一批 n 对直方图、共用一张 d×d 代价矩阵。核按（行, batch 瓦片）分块，瓦片大小 8，矩阵行只读一次、同时乘到 8 个缩放向量上。对数域 log-sum-exp 改成单遍：一边维护当前最大值一边重标定求和，全零边际把势钉在 −∞，不会产出 NaN。迭代超过 100 次时，把 25 步打成 CUDA graph 回放。反向不穿过迭代：前向把对偶势存下来，包络定理保证它们就是对两个边际的精确梯度，反向只是一次广播乘法。
正确性靠差分测试钉死。78 个测试函数覆盖每个导出算子，逐元素对 SciPy（形态学和距离）或 POT（传输），包括 2D / 3D / 更高维、batch、非连续布局、各边界模式、非对称结构元素。传输计划必须还原两个边际，返回的最近背景下标必须真的指到背景。解析梯度对过有限差分方向导数。自托管 CI 带物理 CUDA 卡，每次推送跑的是真核，不是 mock。
测试机是单卡 RTX 4090 D（48 GB）、双路 Xeon Gold 6330。SciPy 和 POT 都绑在一个 CPU 核上。加速比是 GPU 对一个 CPU 核，不是对多核并行的 CPU 实现。论文把 B=1 列单独报了，用来把「换了设备」和「吃了 batch」拆开。
数值一致性：
| 算子族 | 对照 | 最大绝对误差 | 相对 ℓ₂ 误差 |
| 二值形态学 | scipy.ndimage | 0 | 0 |
| 灰度形态学 | scipy.ndimage | 4.77×10⁻⁷ | 2.24×10⁻⁸ |
| 精确欧氏 DT | scipy.ndimage | 2.06×10⁻⁷ | 6.77×10⁻⁹ |
| Chamfer DT（棋盘 / 曼哈顿） | scipy.ndimage | 0 | 0 |
| 暴力 ℓ₂ DT | scipy.ndimage | 2.06×10⁻⁷ | 6.71×10⁻⁹ |
| Sinkhorn 距离 | POT | 1.75×10⁻⁶ | 8.82×10⁻⁸ |
浮点算子的最坏绝对误差压在 1.8×10⁻⁶ 以内。论文记录的行为差异只有 NaN 传播，来自 float32 和 --usefastmath。
吞吐按 inputs/ms 报，越高越好。一张 2D 图或一个 3D 体算一个 input：
| 算子 | 尺寸 | SciPy | TM B=1 | TM B=8 | B=8 相对 SciPy |
| 灰度膨胀 | 256² | 0.756 | 10.3 | 111.1 | 147× |
| 灰度膨胀 | 1024² | 0.040 | 13.2 | 45.5 | 1138× |
| 灰度腐蚀 | 256² | 0.740 | 16.1 | 83.3 | 113× |
| 二值腐蚀 | 256² | 0.865 | 6.9 | 55.6 | 64× |
| 精确 EDT | 1024² | 0.011 | 2.1 | 2.9 | 264× |
| 精确 EDT | 64³ | 0.031 | 6.3 | 10.8 | 348× |
1024² 灰度膨胀 B=8 对应摘要里的 1.1×10³ 倍；64³ EDT 对应「最高 350 倍」。小图上 batch 很值钱：256² 灰度膨胀从 B=1 的 10.3 提到 B=8 的 111.1，10.8 倍。大图和大体积已经能喂饱 GPU，EDT 1024² 只从 2.11 提到 2.86（1.4 倍），128³ 几乎钉在 1.4 inputs/ms。二值形态学在 B=8 落到 36–56 inputs/ms，看算子和尺寸。
Sinkhorn 对 32×32 网格（d=1024 bin）上的成对 ℓ₂ 代价：
| 配置 | POT (ms) | TM (ms) | 加速 | 计划相对误差 |
| scaling，100 步 | 23.0 | 1.2 | 18.8× | 9.21×10⁻⁴ |
| scaling，1000 步（CUDA graph） | 229.7 | 10.1 | 22.8× | 4.45×10⁻⁷ |
| log-domain，200 步 | 35.4 | 3.0 | 11.8× | 4.55×10⁻⁵ |
| batch n=16，100 步 | 367.3 | 8.7 | 42.4× | 1.54×10⁻³ |
边界损失、Hausdorff 代理损失每一步都要重算距离场。以前只能在 CPU 上当预处理，现在可以留在训练循环里，对着已经在 GPU 上的 batch 直接打。对已经用 ndimage 做 mask 清理、开闭运算、孔洞填充的人，改动接近换 import。三维、四维体数据不是特例。
它替代不了 Kornia 的可微二维形态学，形态学核目前前向 only。它也不是给 CPU 训练准备的。场景很具体：CUDA 上的分割或医学影像训练循环，要 SciPy 语义的形态学和 EDT，并且能接受不可微的后处理。需要可微传输损失的，Sinkhorn 这条路径带自定义 autograd，对两个边际都可导。
MIT 许可，依赖只有 PyTorch 和 CUDA 工具链。
作者列了三条。形态学和距离核没有 autograd，只有传输模块可微；腐蚀膨胀理论上可以把次梯度路由到 argmin / argmax 位置，但没做。设备绑定 CUDA，形态学的回退就是 SciPy 自己，传输模块可以退回纯 torch CPU。计算用 float32 加 --usefastmath，NaN 传播不保证对齐参考实现。
加速比的对照是单线程 CPU。论文没有和多核 SciPy 打，也没有和 cuCIM 的 GPU 形态学打。声称填补覆盖缺口，吞吐数字只对 SciPy 和 POT。
空间维封顶 8，超过直接拒。没有连通域和重建类算子，作者把它放进未来工作。测试和计时都在一张 RTX 4090 D 上，换更小的卡，launch 开销占比会更大，小 batch 收益会掉。
暴力距离变换的 SciPy 对照本身就是未优化的正确性预言机，那一列的加速比不能当性能主张读。