This is the documentation of torch_kmeans.
PyTorch implementations of KMeans, Soft-KMeans and Constrained-KMeans
torch_kmeans features implementations of the well known k-means algorithm as well as its soft and constrained variants.
All algorithms are completely implemented as PyTorch modules and can be easily incorporated in a PyTorch pipeline or model. Therefore, they support execution on GPU as well as working on (mini-)batches of data. Moreover, they also provide a scikit-learn style interface featuring
model.fit(), model.predict() and model.fit_predict()
Fully implemented in PyTorch.
GPU support like native PyTorch.
PyTorch script JIT compiled for most performance sensitive parts.
- Works with mini-batches of samples:
each instance can have a different number of clusters.
- Constrained Kmeans works with cluster constraints like:
a max number of samples per cluster or,
a maximum weight per cluster, where each sample has an associated weight.
SoftKMeans is a fully differentiable clustering procedure and can readily be used in a PyTorch neural network model which requires backpropagation.
Unit tested against the scikit-learn KMeans implementation.
GPU execution enables very fast computation even for large batch size or very high dimensional feature spaces (see speed comparison)
Simply install from PyPI
pip install torch-kmeans