PaperlayerView exact revision on arXiv

Source-linked AI summary

TorchMorph: CUDA-accelerated Morphological Transforms

Kai Zhao

arXiv:2608.24738v1cs.CV

TL;DR

TorchMorph addresses the mismatch between CPU-only, single-array morphology tools and batched, higher-dimensional PyTorch imaging pipelines. It provides batch-parallel N-dimensional CUDA operators with SciPy-compatible conventions, achieving close numerical agreement while enabling these transforms inside training loops.

  • Problem

    scipy.ndimage requires device transfers, single-threaded CPU computation, and per-image processing for batched, higher-dimensional GPU tensors.

  • Method

    TorchMorph implements 22 batch-parallel CUDA operators for morphology, distance transforms, and entropy-regularised optimal transport through a SciPy-compatible PyTorch API.

  • Results

    TorchMorph matches binary and chamfer SciPy results exactly, while float-valued operators remain within 1.8×10−6 worst-case absolute error and batched execution improves throughput across tested operators.

  • Takeaways & Limitations

    The library makes classical morphological and distance operators available as validated, batch-parallel components that can run inside modern training loops.

  • Takeaways & Limitations

    Morphology and distance kernels are forward-only, require CUDA, and compute in float32 under --use fast math with potentially different NaN propagation.

Abstract

from arXiv · show

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.

1 Introduction

TorchMorph addresses the mismatch between CPU-only, single-array SciPy morphology and GPU-based, batched, multidimensional imaging pipelines. It provides broad native CUDA coverage while preserving SciPy-compatible semantics.

  • GPU training pipelines must otherwise copy PyTorch tensors to the host, process single arrays on one CPU thread, and copy results back.
  • Three- and four-dimensional workloads are especially costly because morphological sweeps scale with the product of all spatial extents.
  • Existing GPU and vision libraries provide incomplete coverage of N-dimensional operators, SciPy border modes, exact Euclidean distance transforms, or native torch.Tensor execution.
  • TorchMorph offers 22 public operators for morphology, distance transforms, and entropy-regularised optimal transport as batch-parallel CUDA kernels.
  • Its image operators accept (B, C, Spatial...) CUDA tensors with up to eight spatial dimensions and mirror SciPy names, arguments, defaults, and boundary behavior.

2.1 Covered transforms

TorchMorph covers core morphology, distance-transform, and transport workloads used in image processing and learning pipelines. Its operator families combine reusable CUDA primitives with broader host-side compositions.

  • The 22 exported operators are organised into four families plus a structuring-element utility group, with only bold entries backed by dedicated CUDA kernels.
  • Binary and greyscale morphology: Binary and greyscale morphology support mask cleanup, shape decomposition, structure extraction, segmentation post-processing, and label preparation on batched GPU tensors.
  • Distance transforms: The exact Euclidean distance transform uses a separable lower-envelope algorithm, while chamfer transforms support chessboard and taxicab metrics.
  • Entropy-regularised optimal transport: TorchMorph provides batched Sinkhorn optimal transport in scaling and log domains for geometry-aware comparisons of histograms, point clouds, and attention-style matching problems.

2.2 Architecture

TorchMorph separates SciPy-compatible API handling from compact bindings and runtime-rank CUDA kernels. This architecture supports batched tensors while keeping derived operators out of device code.

  • All image operators accept CUDA tensors shaped (B, C, Spatial...) with one to eight spatial dimensions and support SciPy-style arguments and pre-allocated outputs.
  • The porting pattern changes a per-volume SciPy call into one TorchMorph call over the CUDA batch.
  • Python layer: The Python layer resolves structuring elements, validates origins, maps border modes, and composes derived operators.
  • Layered architecture: Figure 1 distinguishes dedicated CUDA kernels from host-side compositions, while Figure 2 separates Python normalization, bindings, and CUDA implementation layers.
  • Binding and kernel layers: The binding layer exposes eight kernel entry points from six CUDA translation units, and runtime spatial rank is bounded at eight so coordinate scratch space remains in registers.

2.3 Kernel design

TorchMorph’s kernels reduce multidimensional indexing and memory overhead through host-precomputed geometry, interior fast paths, separable scans, and batch-aware Sinkhorn execution. Specialized launch and reduction strategies address morphology, distance transforms, and transport separately.

  • Fused morphology kernel: A naive N-dimensional morphology kernel repeatedly maps coordinates and checks boundaries for every output element and structuring-element position.
  • Fused morphology kernel: Host preprocessing flattens active structuring-element entries and stores per-axis and precomputed flat offsets, excluding inactive positions from device execution.
  • Fused morphology kernel: Interior threads use direct flat-offset addressing without boundary tests, while face-near threads implement SciPy-compatible border modes.
  • Fused morphology kernel: Erosion and dilation share a templated reduction kernel, with early termination for binary results once the outcome is decided.
  • Distance transforms: The exact distance transform assigns one block per scanline, builds lower envelopes sequentially, and parallelizes loading and output queries.
  • Distance transforms: Distance-transform paths include 2-D fusion, contiguous-memory transposition, and global-memory fallback for scanlines exceeding shared-memory capacity.
  • Distance transforms: The brute-force transform uses tiled background coordinates and is reserved for validation and anisotropic chamfer cases unavailable through the chamfer entry point.
  • Batch-tiled Sinkhorn solver: Sinkhorn batches share one cost matrix, enabling matrix-row reuse across eight scaling vectors and reduced matrix traffic.

2.4 Testing and reference alignment

TorchMorph validates SciPy-compatible behavior through differential testing across diverse operators, layouts, ranks, and execution paths, with contract and runtime tests guarding API and device semantics.

  • 78 test functions compare exported operators elementwise against SciPy or POT across ranks, layouts, batching, sampling, border modes, and shifted origins.The GPU receives the same per-sample array used by the reference, preventing batching from masking item-level discrepancies.
  • Cross-implementation tests require fused kernels, pure-Torch fallbacks, float32 and float64, CPU and GPU, graph replay, and plain launches to coincide.
  • Contract tests reject unsupported ranks, mismatched shapes, invalid modes or metrics, and non-CUDA input.
  • Runtime tests verify stream and device behavior by issuing work on a side stream and non-default device.

3 Results

TorchMorph is evaluated for numerical agreement and throughput against single-threaded CPU references and POT. Results show exact agreement for binary and chamfer operators, float32-level errors for other float-valued operators, and strong batching gains that vary by transform cost.

  • Numerical agreement: Binary morphology and chamfer distance transforms match SciPy exactly, while float-valued operators stay below 1.8×10−6 worst-case absolute error.The corresponding worst-case relative ℓ2 error is below 9×10−8; NaN propagation is the documented behavioral difference.
  • Throughput: 10.31 to 111.11 inputs/ms lifts grey dilation throughput 10.8× from batch size B = 1 to B = 8.Grey erosion rises from 16.13 to 83.33 inputs/ms, a 5.2× gain, over the same batch-size change.
  • Throughput: 36–56 inputs/ms is reached by binary morphology at B = 8, depending on operator and image size.
  • Throughput: 1.4× is the EDT speedup on 1024^2 images from B = 1 to B = 8, while 128^3-volume EDT remains approximately 1.4 inputs/ms across batch sizes.Heavier transforms already expose enough parallelism in a single large input to approach device saturation; BFDT rises only from 1.15 to 1.25 inputs/ms.
  • Comparison scope: Speed-ups compare GPU execution with one CPU core rather than a well-parallelised CPU implementation.BFDT timings are explicitly treated as correctness-oracle costs, not as a benchmark against an optimized CPU distance-transform implementation.

4 Conclusion

TorchMorph supplies validated, batch-parallel N-dimensional CUDA implementations behind a SciPy-mirroring API, making these operators available within modern tensor workflows. Its scope remains bounded by forward-only morphology and distance kernels, CUDA requirements, and float32 fast-math behavior.

  • 22 operators cover binary and greyscale morphology, exact and approximate distance transforms, and entropy-regularised optimal transport on batch-parallel N-dimensional CUDA tensors.The API mirrors scipy.ndimage, and the contribution is positioned as availability rather than algorithmic novelty.
  • Only the transport module is differentiable; morphology and distance kernels are forward-only, although erosion and dilation admit arg-min/arg-max subgradients.
  • CUDA is required for morphology and distance kernels, with SciPy as the intended fallback; transport falls back to pure Torch on CPU.
  • Float32 computation under --use fast math means NaN propagation is not guaranteed to match the reference.Future work includes autograd for morphological operators and additional connected-component and reconstruction operators.
Loading 2608.24738v1…