Torch Dimensionality Reduction#

torchdr logo

TorchDR is a high-performance dimensionality reduction library built on PyTorch. It provides GPU and multi-GPU accelerated DR methods in a unified framework with a simple, scikit-learn-compatible API.

Key Features#

Feature	Description
High Performance	Engineered for speed with GPU acceleration, `torch.compile` support, and optimized algorithms leveraging sparsity and negative sampling.
Multi-GPU Support	Scale to massive datasets with built-in distributed computing. Use the `torchdr` CLI or `torchrun` for easy multi-GPU execution of compatible modules and methods.
Modular by Design	Every component is designed to be easily customized, extended, or replaced to fit your specific needs.
Memory-Efficient	Natively handles sparsity and memory-efficient symbolic operations. Supports PyTorch DataLoader for streaming large datasets.
Seamless Integration	Fully compatible with the scikit-learn and PyTorch ecosystems. Use familiar APIs and integrate effortlessly into your existing workflows.
Minimal Dependencies	Requires only PyTorch, NumPy, and scikit‑learn; optionally add Faiss for fast k‑NN or KeOps for symbolic computation.

Benchmarks#

TorchDR scales UMAP across 1–8 NVIDIA B200 GPUs on Tahoe-100M, with one-GPU cuML and 64-core umap-learn as references using the same PCA initialization.

TorchDR UMAP runtime scaling from one to eight GPUs on the Tahoe single-cell dataset

Getting Started#

TorchDR offers a user-friendly API similar to scikit-learn where dimensionality reduction modules can be called with the fit_transform method. It seamlessly accepts both NumPy arrays and PyTorch tensors as input, ensuring that the output matches the type and backend of the input.

from sklearn.datasets import fetch_openml
from torchdr import UMAP

x = fetch_openml("mnist_784").data.astype("float32")

z = UMAP(n_neighbors=30).fit_transform(x)

GPU Acceleration: Set device="cuda" to run on GPU. By default (device="auto"), TorchDR uses the input data’s device.

z = UMAP(n_neighbors=30, device="cuda").fit_transform(x)

Multi-GPU: Use the torchdr CLI to parallelize across GPUs with no code changes:

torchdr my_script.py            # Use all available GPUs
torchdr --gpus 4 my_script.py   # Use 4 GPUs

torch.compile: Enable compile=True for additional speed on PyTorch 2.0+.

Backends: The backend parameter controls k-NN and memory-efficient computations:

Backend	Description
`"faiss"`	Fast approximate k-NN via Faiss (Recommended)
`"keops"`	Exact symbolic computation via KeOps with linear memory
`None`	Raw PyTorch

DataLoader for Large Datasets: Pass a PyTorch DataLoader instead of a tensor to stream data batch-by-batch. Requires backend="faiss".

from torch.utils.data import DataLoader, TensorDataset

dataloader = DataLoader(TensorDataset(X), batch_size=10000, shuffle=False)
z = UMAP(backend="faiss").fit_transform(dataloader)

Methods#

Neighbor Embedding#

TorchDR provides a suite of neighbor embedding methods, optimal for data visualization.

Method	Complexity	Multi-GPU	Paper
`UMAP`	O(n)	✅	↗
`LargeVis`	O(n)	✅	↗
`InfoTSNE`	O(n)	✅	↗
`PACMAP`	O(n)	❌	↗
`SNE`	O(n²)	✅	↗
`TSNE`	O(n²)	✅	↗
`TSNEkhorn`	O(n²)	❌	↗
`COSNE`	O(n²)	✅	↗

Note: Quadratic methods support backend="keops" for exact computation with linear memory usage.

Spectral Embedding#

TorchDR provides various spectral embedding methods: PCA, IncrementalPCA, ExactIncrementalPCA, KernelPCA, PHATE. PCA and ExactIncrementalPCA support multi-GPU distributed training via the distributed="auto" parameter.

Examples#

See the examples folder for all examples.

Tahoe-100M. TorchDR UMAP embedding of 80 million single cells across 50 cell lines.

TorchDR UMAP embedding of 80 million cells from the Tahoe single-cell dataset

CIFAR100. (Code) Visualizing the CIFAR100 dataset using DINO features and TSNE.

TSNE on CIFAR100 DINO features

Advanced Features#

Affinities#

TorchDR features a wide range of affinities which can then be used as a building block for DR algorithms. It includes:

Affinities based on k-NN normalizations: SelfTuningAffinity, MAGICAffinity, UMAPAffinity, PHATEAffinity, PACMAPAffinity.
Doubly stochastic affinities: SinkhornAffinity, DoublyStochasticQuadraticAffinity.
Adaptive affinities with entropy control: EntropicAffinity, SymmetricEntropicAffinity.

Evaluation Metrics#

TorchDR provides efficient GPU-compatible evaluation metrics: silhouette_score, knn_label_accuracy, neighborhood_preservation, kmeans_ari.

Installation#

Install the core torchdr library from PyPI:

pip install torchdr  # or: uv pip install torchdr

Note: torchdr does not install faiss-gpu or pykeops by default. You need to install them separately to use the corresponding backends.

Faiss (Recommended): For the fastest k-NN computations, install Faiss. Please follow their official installation guide. A common method is using conda:
```
conda install -c pytorch -c nvidia faiss-gpu
```
KeOps: For memory-efficient symbolic computations, install PyKeOps.
```
pip install pykeops
```

Installation from Source#

If you want to use the latest, unreleased version of torchdr, you can install it directly from GitHub:

pip install git+https://github.com/torchdr/torchdr

Finding Help#

If you have any questions or suggestions, feel free to open an issue on the issue tracker or contact Hugues Van Assel directly.