Docs · Machine learning

kmeans: k-means clustering.

Lloyd's k-means with caller-supplied initial centroids, the equivalent of Spark MLlib's KMeans. Centroids stay resident on the device; the rows are paged. Fixed iteration count, no early stop.

Module
ml_kmeans
Header
clio_ml/kmeans.h
Rows
features only

What it does

Each iteration assigns every row to its nearest centroid and recomputes the centroids from the sums. Partials are reduced per block in a fixed order, so the result is bit-identical whether the cache holds all pages or a fraction of them. k-means parallel initialisation is not implemented: you pass the initial centroids.

Rows

Every column is a feature. clio_ml.COLUMNS["kmeans"] is 0.

Python

python
import numpy as np, clio_ml as cm
X = np.random.default_rng(0).normal(size=(4096, 8)).astype(np.float32)

r = cm.fit_kmeans(X, k=4, iters=5)            # init_centroids defaults to X[:k]
centroids = np.asarray(r.centroids).reshape(4, 8)
print(r.counts)                               # points per cluster, last iteration

cm.init()
ds = cm.dataset(X, cache_pages=8)
r2 = cm.fit_kmeans(ds, k=4, iters=5, init_centroids=X[:4])   # required for a Dataset
cm.finalize()
  • fit_kmeans(X, init_centroids=None, **options). For an in-memory array the default initial centroids are the first k rows; for a Dataset you must pass init_centroids (shape k x features).
  • Options: k (8), iters (10), threads (256).
  • Result KMeansResult: centroids (k * dims, row-major), counts (cluster sizes after the last iteration), stats (faults, puts, evicts, get_errors, put_errors; paged runs only).
  • No one-row predictor is exposed through predict.

C++

cpp
struct KMeansOptions {
  u32 k = 8;
  u32 iters = 10;
  u32 threads = 256;        // <= GV_LB_THREADS
  ClusterExchange cluster;  // default: single node
};
struct KMeansResult {
  std::vector<float> centroids;  // k*dims, row-major
  std::vector<u32> counts;       // k
  VecStats stats;                // paged runs only
};
KMeansResult FitKMeans(Dataset &ds, const std::vector<float> &init_centroids,
                       const KMeansOptions &opt);
KMeansResult FitKMeansPlain(const float *data, u64 n_rows, u32 dims,
                            u64 elems_per_page, u32 nblocks,
                            const std::vector<float> &init_centroids,
                            const KMeansOptions &opt);
std::vector<float> FitKMeansReference(const float *data, u64 n_rows, u32 dims,
                                      u64 elems_per_page, u32 nblocks,
                                      const std::vector<float> &init_centroids,
                                      u32 k, u32 iters);

FitKMeans runs over a paged Dataset and needs the runtime. FitKMeansPlain runs the same kernel over a plain device copy of the data and needs a device but no runtime. FitKMeansReference is the CPU oracle with the same summation order.

ChiMod

Module ml_kmeans, library clio_ml_kmeans_chimod, client clio_ml::chimod::KMeansClient. Compute params are KMeansParams(k, iters, threads, init_centroids). See running as ChiMods.

Related: bisecting k-means, Gaussian mixture, fuzzy c-means.