TorchMorph: CUDA-accelerated Morphological Transforms
Kai Zhao
cs.CV
2026-08-25
TorchMorph ships 22 fused CUDA operators matching scipy.ndimage on batched GPU tensors. Greyscale morphology reaches 1100× single-thread SciPy, exact EDT 350×, Sinkhorn 42× vs POT.
Morphological operators sit in the daily path of mask cleanup, skeletonisation, and boundary-aware losses. In Python the canonical semantics belong to scipy.ndimage: border modes, structuring-element origins, connectivity. Downstream libraries are expected to reproduce that contract. The implementation itself still assumes a single NumPy array in host memory.
A training step does not look like that. Tensors already live on the GPU, and they arrive in batches. Every ndimage call copies to the host, runs single-threaded, and copies back. Existing GPU vision stacks cover a thin slice of the same surface. Kornia offers differentiable 2-D morphology. N-D operators, the full border-mode matrix, and an exact Euclidean distance transform sit outside its coverage. cuCIM talks CuPy, not native torch.Tensor. MONAI sends several morphological post-steps back to SciPy or CuPy. Entropic optimal transport lives in yet another package. A pipeline that wants batched GPU morphology, an exact EDT, and a differentiable transport loss currently glues three libraries with three tensor conventions.
TorchMorph, from Kai Zhao at Shanghai University, puts those operators in one PyTorch extension. The paper frames the contribution as availability, not a new algorithm.
Twenty-two public operators sit in four families: binary morphology, greyscale morphology, distance transforms, and entropy-regularised optimal transport. Dedicated CUDA kernels back only the primitives (erosion, dilation, exact EDT, chamfer, brute-force distance, Sinkhorn). Openings, closings, top-hats, gradients and hole filling are host-side compositions, so a new top-hat costs no device code.
The API copies scipy.ndimage argument for argument: size / footprint / structure, mode, cval, origin, pre-allocated outputs. iterations < 1 means iterate until the result stops changing, which is how binary propagation and fill-holes are built. Inputs are (B, C, Spatial...) CUDA tensors with spatial rank 1 to 8. Porting is a change of import.
Three layers keep SciPy conventions out of the device code. The Python layer normalises arguments and is the only place that knows the parameter contract; structuring elements resolve as structure > footprint > size. A single pybind11 module exposes eight kernel entry points compiled from six CUDA translation units. The kernel layer resolves geometry on the host whenever it can, and writes against a runtime spatial rank with a compile-time cap so coordinate scratch stays in registers.
A naive morphology kernel remaps N-D coordinates and tests every axis at every structuring-element offset. A 3³ element in 3-D is 27 mapped coordinates per voxel. TorchMorph flattens the element on the host into active offsets plus a precomputed flat stride offset; inactive footprint positions never reach the device. Each thread first tests whether it is interior. Interior threads, the bulk of the grid, add the flat offset to the linear index with no per-axis arithmetic and no bounds check. Only threads within one element radius of a face take the general path. The greyscale kernel implements all five SciPy border modes (constant, reflect, nearest, mirror, wrap); the binary kernel uses a single border value. Binary threads can exit early: false decides erosion, true decides dilation, which is cheap on sparse masks.
The exact EDT uses the separable lower-envelope algorithm of Felzenszwalb and Huttenlocher. The 1-D sweep is sequential, so one thread block owns one scanline: a single thread builds the envelope in shared memory, the block cooperates on loads and on the query phase. Parallelism comes from the number of scanlines, B·C·n^{d-1} per pass. A 2-D specialisation for extents up to 2048 keeps the row pass contiguous and fuses the final square root into the column pass. The general path transposes the active axis innermost. Scanlines that exceed the shared-memory budget spill the envelope stack to a lazily allocated global buffer. Chamfer runs separable forward and backward sweeps, with extra diagonal passes for the chessboard metric. The brute-force kernel tiles background coordinates through shared memory in chunks of 256; it is an exactness oracle, and the only path that accepts anisotropic sampling for chamfer metrics.
The Sinkhorn solver takes a batch of n histogram pairs that share one d×d cost matrix. One block owns a (row, batch-tile) pair with a tile of eight items, streams the matrix row once, and applies it to eight scaling vectors in registers. The log-domain update keeps a running maximum and a rescaled sum in a single pass; an all-zero marginal pins the potential at −∞ instead of emitting NaN. From 100 requested iterations upward, chunks of 25 steps are captured into a CUDA graph and replayed. Gradients do not flow through the iteration. Centered dual potentials are stashed in the forward pass; by the envelope theorem they are the exact gradients of the entropic cost with respect to both marginals, so backward is a broadcast multiply.
Correctness is differential. Seventy-eight tests compare every exported operator elementwise against scipy.ndimage or POT, across 2-D, 3-D and higher rank, batches, non-contiguous layouts, every border mode, and shifted origins. Transport plans must recover both marginals. Returned indices must point at a real nearest background pixel. Analytic gradients are checked against a finite-difference directional derivative. CI runs the real kernels on a physical CUDA device.
All timings come from one machine: an RTX 4090 D (48 GB) against a single CPU core of a dual Xeon Gold 6330. SciPy and POT are single-threaded. The speed-ups are GPU versus one core, not versus a well-parallelised CPU implementation. A B=1 column is reported so device effects can be separated from batching.
Numerical agreement:
| Family | Reference | Max abs. err. | Rel. ℓ₂ err. |
| Binary morphology | scipy.ndimage | 0 | 0 |
| Greyscale morphology | scipy.ndimage | 4.77×10⁻⁷ | 2.24×10⁻⁸ |
| Exact Euclidean DT | scipy.ndimage | 2.06×10⁻⁷ | 6.77×10⁻⁹ |
| Chamfer DT (chessboard / taxicab) | scipy.ndimage | 0 | 0 |
| Brute-force ℓ₂ DT | scipy.ndimage | 2.06×10⁻⁷ | 6.71×10⁻⁹ |
| Sinkhorn distance | POT | 1.75×10⁻⁶ | 8.82×10⁻⁸ |
Float-valued operators stay within 1.8×10⁻⁶ absolute error. The only documented behavioural difference is NaN propagation, from float32 plus --usefastmath.
Throughput in inputs/ms (one 2-D image or one 3-D volume per input):
| Operator | Size | SciPy | TM B=1 | TM B=8 | B=8 vs SciPy |
| Grey dilation | 256² | 0.756 | 10.3 | 111.1 | 147× |
| Grey dilation | 1024² | 0.040 | 13.2 | 45.5 | 1138× |
| Grey erosion | 256² | 0.740 | 16.1 | 83.3 | 113× |
| Binary erosion | 256² | 0.865 | 6.9 | 55.6 | 64× |
| Exact EDT | 1024² | 0.011 | 2.1 | 2.9 | 264× |
| Exact EDT | 64³ | 0.031 | 6.3 | 10.8 | 348× |
The 1024² greyscale-dilation number is the 1.1×10³× claim in the abstract; 64³ EDT is the 350× claim. Batching pays on small inputs: grey dilation at 256² climbs from 10.3 inputs/ms at B=1 to 111.1 at B=8, a 10.8× gain. Large inputs already saturate the device. EDT at 1024² only moves from 2.11 to 2.86 (1.4×); at 128³ it sits near 1.4 inputs/ms regardless of batch size. Binary morphology lands in the 36–56 inputs/ms band at B=8.
Sinkhorn on a 32×32 grid (d=1024 bins, pairwise ℓ₂ cost):
| Setup | POT (ms) | TM (ms) | Speed-up | Plan rel. err. |
| scaling, 100 it. | 23.0 | 1.2 | 18.8× | 9.21×10⁻⁴ |
| scaling, 1000 it. (CUDA graph) | 229.7 | 10.1 | 22.8× | 4.45×10⁻⁷ |
| log-domain, 200 it. | 35.4 | 3.0 | 11.8× | 4.55×10⁻⁵ |
| batch n=16, 100 it. | 367.3 | 8.7 | 42.4× | 1.54×10⁻³ |
Boundary losses and Hausdorff surrogates recompute a distance field every training step. That used to be host-side preprocessing. It can now sit inside the loop, on the batch that already lives on the device. Anyone already calling ndimage for mask cleanup, openings, or hole filling is looking at a change of import. 3-D and 4-D volumes are ordinary inputs, not special cases.
This does not replace Kornia's differentiable 2-D morphology: the morphology kernels are forward-only. It is also not a CPU library. The intended setting is a CUDA segmentation or medical-imaging training loop that wants SciPy-faithful morphology and EDT and can live with non-differentiable post-processing. For a geometry-aware loss, the Sinkhorn path is differentiable in both marginals via custom autograd.
MIT licence. Dependencies are PyTorch and a CUDA toolchain.
Three limits are stated. Morphology and distance kernels have no autograd; only transport is differentiable. Erosion and dilation admit subgradients routed to the arg-min / arg-max site, but that path is not implemented. The morphology kernels require CUDA, with SciPy as the intended fallback; transport can drop back to pure torch on CPU. Arithmetic is float32 under --usefastmath, and NaN propagation is not guaranteed to match the reference.
The speed-ups compare a GPU with one CPU core. There is no multi-core SciPy baseline and no head-to-head against cuCIM's GPU morphology. Coverage is argued against a table of libraries; throughput is argued only against SciPy and POT.
Spatial rank is capped at eight. Connected-component and reconstruction operators are absent and listed as future work. Every number comes from a single RTX 4090 D; on a smaller card, launch overhead would eat more of the B=1 column. The brute-force SciPy timings are a correctness oracle, not a performance claim.