Skip to main content

PyTorch and Spectral Clustering

PyTorch 2.14.1 provides tensor, distance, and eigendecomposition operations, but it does not provide a maintained high-level spectral-clustering estimator. A short handwritten implementation can easily mix Laplacian conventions or choose the wrong eigenvectors. This example uses PyTorch at the tensor boundary and scikit-learn for the complete estimator.

1. Generate deterministic data​

import matplotlib.pyplot as plt
import torch
from sklearn.cluster import SpectralClustering
from sklearn.datasets import make_moons
from sklearn.metrics import adjusted_rand_score

X, reference = make_moons(n_samples=240, noise=0.05, random_state=42)
X_tensor = torch.tensor(X, dtype=torch.float64)

2. Construct and inspect the affinity​

torch.cdist returns Euclidean distances. Squaring before applying the RBF formula matches exp⁡(−γ∥xi−xj∥2)\exp(-\gamma\lVert x_i-x_j\rVert^2).

gamma = 15.0
squared_distances = torch.cdist(X_tensor, X_tensor).square()
affinity_tensor = torch.exp(-gamma * squared_distances)
affinity_tensor.fill_diagonal_(1.0)
assert torch.allclose(affinity_tensor, affinity_tensor.T)
assert torch.all(affinity_tensor >= 0)
affinity = affinity_tensor.numpy()

3. Fit through the maintained estimator​

model = SpectralClustering(
    n_clusters=2,
    affinity="precomputed",
    assign_labels="cluster_qr",
    random_state=42,
)
labels = model.fit_predict(affinity)
assert adjusted_rand_score(reference, labels) > 0.95

4. Plot the partition​

plt.scatter(X[:, 0], X[:, 1], c=labels, cmap="viridis", s=24)
plt.title("Spectral clustering from a PyTorch-built affinity")
plt.xlabel("Feature 1")
plt.ylabel("Feature 2")
plt.close()

This boundary keeps tensor preprocessing in a PyTorch workflow while retaining scikit-learn's tested graph, eigensolver, and label-assignment behavior. The dense affinity still has quadratic memory cost and transfers data to CPU NumPy.