minibatch-kmeans 0.1.1

Mini-batch k-means clustering implemented in PyTorch

minibatch-kmeans

This package implements mini-batch k-means following Sculley (WWW 2010) (Web-Scale K-Means Clustering). scikit-learn's MiniBatchKMeans is based on the same algorithm but is CPU-only; this library is for PyTorch pipelines that want native tensors and GPU acceleration.

Requirements

• Python 3.10+ (see pyproject.toml)
• PyTorch 2.x (torch>=2.0; use pytorch.org for the wheel that matches your platform / CUDA).

Installation

pip install minibatch-kmeans

Quick start

import torch
from minibatch_kmeans import MiniBatchKMeans

X = torch.randn(5000, 4, dtype=torch.float32)
km = MiniBatchKMeans(n_clusters=8, dtype=torch.float32)
km.fit(X, batch_size=256, max_iter=20)
labels = km.predict(X)

Incremental training

km = MiniBatchKMeans(n_clusters=8, dtype=torch.float32)
for batch in X.split(256):
    km.partial_fit(batch)
labels = km.predict(X)

Example notebook: notebooks/partial_fit_demo.ipynb

Documentation

API reference: mb-14.github.io/minibatch-kmeans

Benchmarks & evaluation

See evaluation/README.md for benchmarks and evaluation.

Development

pip install -e ".[dev]"
python -m unittest discover -s tests -v

Generate local API docs (HTML in pdoc_html/):

bash scripts/gen_api_docs.sh

Related projects

• AnswerDotAI/fastkmeans
• svg-project/flash-kmeans
• Faiss wiki
• rapidsai/cuml