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.

Governance and Releases#

TorchDR is community-maintained; see the governance policy and contribution guide. Releases follow Semantic Versioning, generally target two feature releases per year, and ship when tests—including the latest two stable PyTorch releases—pass, documentation and release notes are complete, packages build, and maintainers approve. See the release process and release notes.