Abstract
TorchMorph is a PyTorch extension providing GPU-accelerated morphological and distance-transform operators across up to eight dimensions with a SciPy-compatible API.
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.
Community
This paper presents TorchMorph, a CUDA-accelerated morphological transforms library for PyTorch. It exposes 22 operators (binary and greyscale morphology, exact/chamfer/brute-force distance transforms, entropic optimal transport) as fused CUDA kernels running directly on (B, C, Spatial...) tensors with up to 8 spatial dimensions, behind an API that mirrors scipy.ndimage argument-for-argument, so porting is a change of import. Existing GPU alternatives don't cover this: Kornia is 2-D only, cuCIM runs on CuPy and caps its exact EDT at 3-D, MONAI forwards to SciPy channel-wise.
Against single-threaded references it reaches up to 1.1×10³× scipy.ndimage throughput on greyscale morphology, 350× on exact EDT, and 42× over POT on Sinkhorn. Binary and chamfer operators match SciPy bit-exactly; float-valued operators agree to within 1.8×10⁻⁶. MIT licensed.
Every morphology op I've needed inside a training loop meant either a .cpu() round trip or a hand-rolled kernel I never fully trusted. TorchMorph kills both — 22 operators, batched, GPU-native, drop-in. That's the kind of plumbing that saves actual hours, not the kind that looks good in a benchmark table. The part I'd want to see before adopting it in earnest is how structuring element handling scales to 3D volumes, because that's where naive implementations usually fall apart. For 2D mask pipelines though, this is a straight upgrade.
This is an automated message from the Librarian Bot. I found the following papers similar to this paper.
The following papers were recommended by the Semantic Scholar API
- Fast and Memory-Efficient Wavelet Convolutions via I/O-Aware Reformulation (2026)
- FlashBEV: Fast and Memory-Efficient Exact BEV Transformation with IO-Awareness (2026)
- Reducing Data Movement in the Galerkin Product of Block Algebraic Multigrid on GPUs (2026)
- Fast LapSum: Exact Differentiable Top-k at Million Scale (2026)
- Native Multi-Dimensional Subquadratic Operators via Input Dependent Long Convolutions (2026)
- MegaSlide-DiT: Memory-Centric Adaptation and Deformable Local Attention for Efficient Video Diffusion (2026)
- Flash EQ-Linear: Accelerating Equivariant Linear Layers via Group-wise Discrete Fourier Transform (2026)
Please give a thumbs up to this comment if you found it helpful!
If you want recommendations for any Paper on Hugging Face checkout this Space
You can directly ask Librarian Bot for paper recommendations by tagging it in a comment: @librarian-bot recommend
Get this paper in your agent:
hf papers read 2608.24738 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper