dbscan-torch

GPU-accelerated 2D DBSCAN as a native PyTorch op. Tensor in, tensor out, on the same device, with no host to device round-trip. Parallel CPU and CUDA implementations, dispatched automatically on X.device.

Based on the grid-hashing approach from Wang et al. (SIGMOD 2020), reimplemented as a native PyTorch extension with CUDA support.

Install

pip install dbscan-torch

CPU-only wheels work out of the box on Linux and macOS (not tested in Windows). Building with CUDA requires the CUDA toolkit and a CUDA-enabled torch install at build time.

Installing alongside a pinned torch

The C++ extension is built against whatever torch lands in pip's isolated build environment. If your runtime env has torch pinned to a version that differs from "latest in the supported range," the build will compile against the build-env torch and fail to import against your runtime torch (PyTorch does not preserve binary compatibility across minor versions). dbscan-torch detects this on import and tells you exactly what to do. The fix is to build against your existing torch:

pip install --no-build-isolation dbscan-torch

Use

import torch
from dbscan_torch import dbscan2d

X = torch.randn(100_000, 2, device="cuda")
labels = dbscan2d(X, eps=0.5, min_samples=10)
# labels: int32 tensor of shape (N,) on the same device as X.
# -1 = noise, otherwise a cluster id in [0, n_clusters).

X must be a float32 tensor of shape (N, 2). The function dispatches to the CPU or CUDA implementation automatically based on X.device. Thread count for the CPU path follows torch.get_num_threads(), which is the same method that controls parallelism for any other torch op.

See examples/ for runnable scripts.

Why this exists

sklearn.cluster.DBSCAN is single-threaded CPU and forces you to copy your tensor off the device, cluster, then copy labels back. dbscan-torch keeps the whole pipeline on the GPU while being magnitudes faster than sklearn and lighter than dbscan-python, while also supporting GPU acceleration for massive gains on large N.

The third row is the reference implementation of Wang et al., "Theoretically-Efficient and Practical Parallel DBSCAN" (SIGMOD 2020), distributed as the dbscan-python package (pip install dbscan).

Scaling

Median time per call against N across all five implementations, on synthetic 2D Gaussian clusters at constant density. CPU runs (sklearn, Wang et al., dbscan-torch CPU) were measured locally on a MacBook Pro Intel i9. GPU runs (cuml.DBSCAN, dbscan-torch CUDA) were measured on a Modal T4 instance.

sklearn goes off-scale past 100k. cuml.DBSCAN is included for reference but scales poorly on 2D inputs in this benchmark and caps out by 500k. The CUDA line stays under 100ms across the whole range while the CPU implementations scale linearly. At 10M points dbscan-torch on CUDA is ~23x faster than Wang et al. and ~11x faster than the torch CPU path.

Benchmark Details

(- means run skipped. cuml caps out by 500k in this benchmark.)

Author

Ryan Peters

License

MIT. See LICENSE.