TorchMorph：CUDA加速によるモルフォロジー変換

要旨

モルフォロジ変換は、形状およびマスク処理のための長年にわたるツールですが、Pythonエコシステムにおける事実上の参照実装であるscipy.ndimageはCPU専用かつ単一配列専用であり、高コストなデバイスからホストへの往復を伴わない限りGPUトレーニングループ内では使用できません。PyTorch上に構築されたGPUビジョンライブラリは、これらのオペレータの狭いサブセットしかカバーしておらず、通常は2次元の空間次元と平坦な構造要素に制限されています。我々は、このギャップを埋める軽量なPyTorch拡張であるTorchMorphを提案します。TorchMorphは、二値モルフォロジ、グレースケールモルフォロジ、厳密および近似距離変換、エントロピー正則化付き最適輸送をカバーする22の公開オペレータを提供します。これらはすべて融合CUDAカーネルとして実装され、最大8つの空間次元を持つ (B, C, Spatial...) 形式のCUDAテンソルに対して直接動作します。APIは意図的にscipy.ndimageと引数単位で一致しており、境界モード、構造要素の原点、事前割り当て出力などを含むため、既存のパイプラインはインポートを変更するだけで移植できます。本論文では、各オペレータファミリーの背後にある階層的アーキテクチャとカーネル設計について説明します。シングルスレッドのCPU参照と比較すると、バッチ実行はグレースケールモルフォロジでscipy.ndimageの最大1.1e3倍、厳密なユークリッド距離変換で最大350倍のスループットに達し、シンクホーンソルバーは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.