TorchMorph：CUDA 加速的形態學轉換

摘要

形態學變換長期以來是形狀與遮罩處理的重要工具，但 Python 生態系統中事實上的參考實作 scipy.ndimage 僅限 CPU、僅支援單一陣列，因此若無昂貴的裝置到主機（device-to-host）往返，便無法在 GPU 訓練迴圈中使用。建基於 PyTorch 的 GPU 視覺函式庫僅涵蓋這些運算子的一小部分，且通常受限於二維空間與平面結構元素。我們提出 TorchMorph，這是一個輕量級 PyTorch 擴充套件，用以填補此缺口。TorchMorph 提供 22 個公開運算子，涵蓋二元形態學、灰階形態學、精確與近似距離變換，以及熵正則化最佳傳輸；這些運算子均以融合的 CUDA 內核實作，可直接操作 (B, C, Spatial...) 形式的 CUDA 張量，最多支援八個空間維度。其 API 特意與 scipy.ndimage 逐參數對應，包括邊界模式、結構元素原點以及預先分配的輸出，因此現有流程只需更改 import 即可移植。我們描述了分層架構以及每個運算子系列背後的內核設計。與單執行緒 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.