torch-kmeans 0.2.0

Highlights

• Fully implemented in PyTorch. (PyTorch and Numpy are the only package dependencies!)

• 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)

Installation

Simply install from PyPI

pip install torch-kmeans

Usage

Pytorch style usage

import torch
from torch_kmeans import KMeans

model = KMeans(n_clusters=4)

x = torch.randn((4, 20, 2))   # (BS, N, D)
result = model(x)
print(result.labels)

Scikit-learn style usage

import torch
from torch_kmeans import KMeans

model = KMeans(n_clusters=4)

x = torch.randn((4, 20, 2))   # (BS, N, D)
model = model.fit(x)
labels = model.predict(x)
print(labels)

or

import torch
from torch_kmeans import KMeans

model = KMeans(n_clusters=4)

x = torch.randn((4, 20, 2))   # (BS, N, D)
labels = model.fit_predict(x)
print(labels)

Examples

You can find more examples and usage in the detailed example notebooks.