TorchMorph: CUDA-가속 형태학적 변환

초록

형태학적 변환은 형태(shape) 및 마스크 처리를 위한 오랜 기간 사용된 도구이지만, Python 생태계에서 사실상 표준 참조 구현인 scipy.ndimage는 CPU 전용이며 단일 배열만 지원하므로, 값비싼 장치-호스트(device-to-host) 왕복 없이는 GPU 훈련 루프 내에서 사용할 수 없다. PyTorch 기반 GPU 비전 라이브러리들은 이러한 연산자 중 극히 일부만을 다루며, 대개 2차원 공간과 평평한 구조 요소로 제한된다. 우리는 이러한 격차를 해소하는 경량 PyTorch 확장 라이브러리 TorchMorph를 제시한다. TorchMorph는 이진 형태학, 그레이스케일 형태학, 정확 및 근사 거리 변환, 엔트로피 정규화 최적 수송을 포함한 22개의 공개 연산자를 제공하며, 모두 최대 8개의 공간 차원을 가진 (B, C, Spatial...) CUDA 텐서에 직접 연산하는 융합(fused) CUDA 커널로 구현된다. API는 경계 모드, 구조 요소 원점, 사전 할당 출력을 포함하여 scipy.ndimage의 인자 하나하나를 의도적으로 그대로 따르므로, import 변경만으로 기존 파이프라인을 이식할 수 있다. 우리는 각 연산자 계열을 뒷받침하는 계층적 아키텍처와 커널 설계를 설명한다. 단일 스레드 CPU 참조 구현과 비교했을 때, 배치 실행은 그레이스케일 형태학에서 scipy.ndimage 대비 최대 1.1 × 10^3배의 처리량, 정확한 유클리드 거리 변환에서 최대 350배의 처리량에 도달하며, Sinkhorn 솔버는 POT보다 최대 42배 빠르게 실행된다. 이진 및 Chamfer 연산자는 SciPy 대응 연산자와 정확히 일치하고, 모든 실수(float) 값 연산자는 CPU 참조 구현과 최대 1.8 × 10^-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.