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.

Key Features

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.

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:

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.

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.

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

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.