Skip to main content

Agglomerative clustering with PyTorch data

PyTorch 2.14 does not ship a native high-level agglomerative-clustering estimator. torch.cdist can inspect point distances, but a complete hierarchy also needs a defined linkage rule, efficient distance updates, deterministic tie handling, stopping rules, and tests. This example uses scikit-learn and SciPy for those maintained capabilities.

1. Prepare deterministic tensor data​

import numpy as np
import torch
from sklearn.datasets import make_blobs
from sklearn.preprocessing import StandardScaler

torch.manual_seed(42)
X, _ = make_blobs(n_samples=150, centers=3, cluster_std=0.60, random_state=42)
X_scaled = StandardScaler().fit_transform(X).astype("float32")
points = torch.from_numpy(X_scaled)
assert points.shape == (150, 2)

2. Inspect Euclidean distances​

pairwise = torch.cdist(points, points)
assert pairwise.shape == (150, 150)
np.testing.assert_allclose(torch.diag(pairwise).numpy(), 0.0, atol=1e-3)

Materializing all pairwise distances takes quadratic memory. It demonstrates the input geometry but is not itself average, complete, single, or Ward linkage.

3. Fit the maintained implementation​

from sklearn.cluster import AgglomerativeClustering

model = AgglomerativeClustering(n_clusters=3, metric="euclidean", linkage="ward")
labels = model.fit_predict(points.numpy())
assert np.unique(labels).size == 3

Ward is tied to Euclidean distance. For Manhattan, cosine, or a precomputed dissimilarity matrix, choose a compatible linkage such as average or complete and justify that geometry for the domain.

4. Verify the cut and visualize the tree​

import matplotlib.pyplot as plt
from scipy.cluster.hierarchy import dendrogram, fcluster, linkage
from sklearn.metrics import adjusted_rand_score, silhouette_score

linked = linkage(X_scaled, method="ward", metric="euclidean")
cut_labels = fcluster(linked, t=3, criterion="maxclust")
assert adjusted_rand_score(labels, cut_labels) == 1.0
silhouette = silhouette_score(X_scaled, labels)
assert 0.70 < silhouette <= 1.0

fig, ax = plt.subplots(figsize=(9, 4))
dendrogram(linked, truncate_mode="lastp", p=12, show_leaf_counts=True, ax=ax)
ax.set(title="Truncated Ward dendrogram", xlabel="merged node or observation", ylabel="Ward distance")
fig.tight_layout()
plt.close(fig)

For learned PyTorch embeddings, detach them, move them to CPU, and treat clustering as a separate fitted artifact. A dendrogram displays the algorithm's merge history under the chosen representation; it does not prove a uniquely correct taxonomy.