Skip to main content

K-Means with TensorFlow data

TensorFlow 2.21 does not expose a maintained high-level K-Means estimator in its public API. The former tf.compat.v1.estimator.experimental.KMeans path belongs to the removed Estimator stack and is not an appropriate TensorFlow 2 example. In production, use a maintained clustering library such as scikit-learn and convert tensors at the boundary. TensorFlow remains useful for feature computation or for illustrating the assignment and update operations.

1. Prepare deterministic tensor data​

The example uses compact, well-separated Gaussian blobs so that the behavior is easy to inspect. Scaling is fitted once on the observations being explored; in a deployment workflow, fit preprocessing on the reference/training population and reuse it for later data.

import numpy as np
import tensorflow as tf
from sklearn.datasets import make_blobs
from sklearn.preprocessing import StandardScaler

tf.random.set_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 = tf.constant(X_scaled)
assert points.shape == (240, 2)

2. Fit the maintained estimator​

KMeans minimizes Euclidean within-cluster sum of squares. n_init=20 makes the number of restarts explicit instead of relying on a version-dependent default; random_state makes centroid initialization reproducible.

from sklearn.cluster import KMeans

kmeans = KMeans(n_clusters=4, n_init=20, random_state=42)
labels = kmeans.fit_predict(points.numpy())
centers = tf.constant(kmeans.cluster_centers_, dtype=tf.float32)

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

3. Inspect one TensorFlow assignment step​

The following tensor calculation is not another estimator. It verifies the core assignment operation for the fitted centers. Squared Euclidean distance is enough because the square root would not change the nearest center.

squared_distances = tf.reduce_sum(
    tf.square(points[:, tf.newaxis, :] - centers[tf.newaxis, :, :]), axis=2
)
tensor_labels = tf.argmin(squared_distances, axis=1, output_type=tf.int32)
np.testing.assert_array_equal(tensor_labels.numpy(), labels)

A complete implementation would additionally need seeded initialization, repeated restarts, empty-cluster handling, convergence checks, numerical tests, and batching. The short check above intentionally does not claim those capabilities.

4. Evaluate and plot​

Silhouette measures cohesion and separation under the supplied metric. It is useful for comparing plausible configurations on the same representation, but a high value does not establish that clusters are real, stable, or useful.

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​

  • Keep tensors in TensorFlow while a neural model produces embeddings.
  • Convert the fixed embedding matrix to NumPy for scikit-learn K-Means.
  • Persist the preprocessing, embedding model, and K-Means estimator together.
  • Refit deliberately when the population changes; cluster identifiers have no inherent ordering or semantic label.

This boundary uses each library for a maintained capability and avoids presenting a teaching loop as a production clustering API.