Skip to main content

K-Means with PyTorch data

PyTorch 2.14 does not ship a maintained high-level K-Means estimator. It provides tensor operations such as torch.cdist, which are useful for feature pipelines and for understanding the algorithm. For a production estimator with initialization, convergence, and empty-cluster behavior already tested, use scikit-learn and cross the tensor/NumPy boundary explicitly.

1. Prepare deterministic 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=240, centers=4, cluster_std=0.55, random_state=42)
X_scaled = StandardScaler().fit_transform(X).astype("float32")
points = torch.from_numpy(X_scaled)
assert points.shape == (240, 2)

2. Fit K-Means at the library boundary​

from sklearn.cluster import KMeans

kmeans = KMeans(n_clusters=4, n_init=20, random_state=42)
labels = kmeans.fit_predict(points.numpy())
centers = torch.from_numpy(kmeans.cluster_centers_.astype("float32"))

assert np.unique(labels).size == 4
assert kmeans.n_iter_ <= kmeans.max_iter
print(f"Inertia: {kmeans.inertia_:.2f}")

The explicit n_init performs multiple seeded starts and keeps the solution with the lowest inertia. A single random tensor selection is not an equivalent replacement.

3. Verify the assignment operation​

distances = torch.cdist(points, centers, p=2)
tensor_labels = torch.argmin(distances, dim=1)
np.testing.assert_array_equal(tensor_labels.numpy(), labels)

This is one inference step for fixed centers. A complete custom Lloyd implementation would also need k-means++ or another justified initializer, empty-cluster recovery, a tolerance-based stopping condition, multiple restarts, device-aware tests, and careful batching. Autograd is not needed for ordinary K-Means because Lloyd's discrete assignment step is not differentiated.

4. Evaluate and visualize​

import matplotlib.pyplot as plt
from sklearn.metrics import silhouette_score

silhouette = silhouette_score(X_scaled, labels)
assert 0.70 < silhouette <= 1.0

fig, ax = plt.subplots()
ax.scatter(X_scaled[:, 0], X_scaled[:, 1], c=labels, cmap="viridis", s=18)
ax.scatter(centers[:, 0], centers[:, 1], c="red", marker="x", s=100, label="centers")
ax.set(title=f"K-Means (silhouette={silhouette:.2f})", xlabel="scaled feature 1", ylabel="scaled feature 2")
ax.legend()
fig.tight_layout()
plt.close(fig)

5. Practical boundary​

Use PyTorch to compute embeddings or other learned representations, detach those representations from the computation graph, move them to CPU, and fit the maintained clustering estimator. Validate stability across seeds and samples as well as internal metrics; cluster labels alone are not ground truth.