TorchMorph：CUDA加速的形态学变换

摘要

形态学变换是形状与掩膜处理的长期工具，但 Python 生态系统中事实上的参考实现 scipy.ndimage 仅支持 CPU 且仅能处理单一数组，因此在 GPU 训练循环中若不经由昂贵的设备到主机往返传输便无法使用。基于 PyTorch 构建的 GPU 视觉库仅覆盖这些算子中一个狭窄的子集，通常仅限于二维空间和扁平结构元素。我们提出了 TorchMorph，一个填补这一空白的轻量级 PyTorch 扩展。TorchMorph 公开了 22 个公共算子，涵盖二值形态学、灰度形态学、精确与近似距离变换，以及熵正则化最优传输；所有这些算子均以融合 CUDA 内核实现，可直接操作形状为 (B, C, Spatial...) 且至多具有八个空间维度的 CUDA 张量。其 API 刻意逐参数地镜像 scipy.ndimage，包括边界模式、结构元素原点及预分配输出，因此现有流水线只需更改导入语句即可迁移。我们描述了各算子家族背后的分层架构和内核设计。相较于单线程 CPU 参考实现，批处理执行在灰度形态学上的吞吐量可达 scipy.ndimage 的 1.1e3 倍，在精确欧氏距离变换上可达 350 倍；同时，Sinkhorn 求解器比 POT 快至多 42 倍。二值算子和倒角算子可精确复现其 SciPy 对应实现，且每个浮点值算子与 CPU 参考实现的绝对误差均在 1.8e-6 以内。TorchMorph 以 MIT 许可证发布，见 https://intcomp.github.io/tm。

English

Morphological transforms are long-standing tools for shape and mask processing, but the de facto reference implementation in the Python ecosystem, i.e. scipy.ndimage, is CPU-only, single-array, and therefore unusable inside a GPU training loop without an expensive device-to-host round trip. GPU vision libraries built on PyTorch cover a narrow subset of these operators, typically restricted to two spatial dimensions and flat structuring elements. We present TorchMorph, a lightweight PyTorch extension that closes this gap. TorchMorph exposes 22 public operators covering binary morphology, greyscale morphology, exact and approximate distance transforms, and entropy-regularised optimal transport, all implemented as fused CUDA kernels that operate directly on (B, C, Spatial...) CUDA tensors with up to eight spatial dimensions. The API deliberately mirrors scipy.ndimage argument-for-argument, including border modes, structuring-element origins and pre-allocated outputs, so that existing pipelines port with a change of import. We describe the layered architecture and the kernel designs behind each operator family. Against single-threaded CPU references, batched execution reaches up to 1.1e3 times the throughput of scipy.ndimage on greyscale morphology and up to 350x on exact Euclidean distance transforms, while the Sinkhorn solver runs up to 42x faster than POT. Binary and chamfer operators reproduce their SciPy counterparts exactly, and every float-valued operator agrees with the CPU reference to within 1.8e-6 absolute error. TorchMorph is released under the MIT licence at https://intcomp.github.io/tm.