让形态学运算在GPU上高效运行，支持多维张量与训练循环。
TorchMorph: CUDA-accelerated Morphological Transforms
- 基于CUDA融合内核，支持最多八维空间张量的形态学操作
- 批量处理速度比CPU快1100倍，精确距离变换快350倍
- 兼容scipy接口，现有代码改一行import即可迁移
形态学运算是处理形状和掩码的经典工具，但当前Python生态中的标准实现scipy.ndimage仅限于CPU、单数组处理，无法在GPU训练循环中直接使用，需昂贵的设备-主机往返。基于PyTorch的GPU视觉库通常只覆盖二维、扁平结构元素的有限算子。本文提出TorchMorph，一个轻量级PyTorch扩展，填补该空白。TorchMorph提供22个公开算子，涵盖二值形态学、灰度形态学、精确与近似距离变换、熵正则最优传输，全部以融合CUDA内核实现，可直接作用于(B, C, Spatial...)格式的CUDA张量，支持最高八维空间。其API刻意模仿scipy.ndimage，参数一一对应，包含边界模式、结构元素原点和预分配输出，使现有流程仅需修改导入语句即可迁移。我们描述了各算子家族的分层架构与内核设计。相较于单线程CPU参考实现，批量执行时灰度形态学吞吐量提升达1.1×10³倍，精确欧氏距离变换提升350倍，Sinkhorn求解器速度比POT快42倍。二值与chamfer算子与SciPy结果完全一致，所有浮点算子与CPU参考结果绝对误差不超过1.8×10⁻⁶。TorchMorph以MIT许可证发布于https://intcomp.github.io/tm。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。