Torch Dimensionality Reduction

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.

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 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:

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)

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

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

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

See the examples folder for all examples.

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

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

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.

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

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

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

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

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.