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.
Figures & tables
GPU
Torch
Batch
N -D
Sem.
EDT
OT
scipy.ndimage [ 7 ]
✗
✗
✗
✓
✓
✓
✗
scikit-image [ 14 ]
✗
✗
✗
✓
∼
✓
✗
OpenCV [ 13 ]
∼
✗
✗
✗
✗
✓
✗
Kornia [ 10 ]
✓
✓
✓
✗
✗
∼
✗
cuCIM [ 11 ]
✓
✗
✗
✗
∼
✓
✗
MONAI [ 12 ]
✓
✓
∼
✗
✗
∼
✗
Table 1: Coverage of the capabilities TorchMorph targets. ✓ = supported, ∼ = partial or indirect, ✗ = not supported. Torch = operates on torch.Tensor natively; Batch = a single call over a leading batch dimension rather than a Python loop; N -D = more than three spatial dimensions; Sem. = matches scipy.ndimage border modes, structuring-element origins and iteration conventions; EDT = exact Euclidean distance transform.
Figure 3: The three implementation layers. SciPy-compatible argument conventions are confined to the Python layer; the CUDA layer sees only normalised geometry.
Operator family
Reference
Max abs. err.
Rel. ℓ2 err.
Binary morphology
ndimage
0
0
Grey morphology
ndimage
4.77×10−7
2.24×10−8
Euclidean DT
ndimage
2.06×10−7
6.77×10−9
Chamfer DT (chessb.)
ndimage
0
0
Chamfer DT (taxicab)
ndimage
0
0
Brute-force DT ( ℓ2 )
ndimage
2.06×10−7
6.71×10−9
Table 2: Worst-case numerical error over all tested configurations against the CPU reference implementations. Binary morphology and chamfer distance transforms match scipy.ndimage exactly in the tested cases, while float-valued operators differ only at float32-level numerical precision. Sinkhorn distance is compared against POT.
Operator
Size
Throughput (inputs/ms) ↑
SciPy
TM
B=1
B=2
B=4
B=8
2562
0.865
6.9
13.7
27.8
55.6
Binary erosion
10242
0.056
6.8
12.8
22.2
35.7
2562
0.851
6.7
13.5
27.0
52.6
Binary dilation
10242
0.057
6.7
12.5
21.7
35.7
Table 3: Per-input throughput across representative operators, input sizes, and batch sizes. Throughput is reported in inputs/ms (higher is better), where one input denotes one 2-D image or one 3-D volume. TM = TorchMorph ; its values are rounded to one decimal place, while the SciPy column keeps three, since one decimal would collapse most of its entries to 0.0.
Configuration
POT
TM
Speed-up
Rel. err.
(ms)
(ms)
(plan)
scaling, 100 it.
23.0
1.2
18.8×
9.21×10−4
scaling, 1000 it.
229.7
10.1
22.8×
4.45×10−7
log-domain, 200 it.
35.4
3.0
11.8×
4.55×10−5
batch n=16 , 100 it.
367.3
8.7
42.4×
1.54×10−3
Table 4: Entropic optimal transport against POT on 2-D grid problems with d=322 bins (TM = TorchMorph ). The 1000-iteration row uses CUDA-graph replay. Timings are rounded to one decimal place; speed-ups are computed from the unrounded values.
Image morphing aims to produce a smooth and semantically consistent transition between two input images. Existing diffusion-based morphing methods either require expensive per-pair optimization or rely on implicit spatial alignment, which easily fails under large layout discrepancies. To address these limitations, we propose AlignMorph, a novel tuning-free diffusion framework guided by the principle of transport-then-denoise. We explicitly decouple geometric alignment from generative denoising to avoid structural entanglement. Our framework consists of two core components. (1) Global Semantic Transport, which achieves diffusion-compatible semantic alignment via entropic optimal transport and reliability-aware latent warping; and (2) Coordinate-Aligned Generation, which uses a symmetric bi-phase attention handoff to maintain consistent spatial coordinates throughout denoising. Without any tuning, AlignMorph effectively eliminates ghosting and achieves superior structural coherence and temporal smoothness on morphing benchmarks. Code is available at https://github.com/51xOne/Alignmorph.
Wuyi Liu, Xu Han, Yuren Chen +3
Huazhong University of Science and Technology · Beijing Jiaotong University · Beihang University +1
Entropic regularized optimal transport (OT) via the Sinkhorn algorithm has become a fundamental tool in machine learning, yet existing implementations either suffer from numerical instability for small regularization parameters or incur significant overhead from deep learning frameworks. We present FastSinkhorn, a lightweight, native CUDA implementation of the log-domain Sinkhorn algorithm that combines warp-level shuffle reductions with shared-memory tiling to achieve high GPU utilization without sacrificing numerical stability. Our solver operates entirely in the log-domain, enabling robust computation for regularization parameters as small as epsilon = 10^{-4} where standard-domain methods fail. On dense OT problems with n = m = 8192, our implementation achieves 12x speedup over the widely-used POT library and 5.9x speedup over GPU-accelerated PyTorch baselines, while consuming only 256 MB of GPU memory. We validate our solver on image color transfer, 3D point cloud matching, and convergence analysis, demonstrating that native CUDA kernels with careful numerical treatment provide a practical and efficient foundation for large-scale optimal transport computation.
3D point cloud perception remains tightly coupled to custom CUDA operators for spatial operations, limiting portability and efficiency on non-NVIDIA, AMD, and embedded hardware. We introduce PointTransformerX (PTX), a fully PyTorch-native vision transformer backbone for 3D point clouds, removing all custom CUDA operators and external libraries while retaining competitive accuracy. PTX introduces 3D-GS-RoPE, a rotary positional embedding that encodes 3D spatial relationships directly in self-attention without neighborhood construction, and further replaces sparse convolutional patch embedding with a linear projection. PTX explores inference-time scaling of attention windows to improve accuracy without retraining. With a redesigned feed-forward network, PTX achieves 98.7% of PointTransformer V3's accuracy on ScanNet with 79.2% fewer parameters and executing 1.6\times faster while requiring just 253 MB memory. PTX runs natively on NVIDIA GPUs, AMD GPUs (ROCm), and CPUs, providing an efficient and portable foundation for point cloud perception.
Laurenz Reichardt, Nikolas Ebert, Oliver Wasenmüller