Skip to main content

No project description provided

Project description

K Means using PyTorch

PyTorch implementation of kmeans for utilizing GPU

Getting Started

import torch
import numpy as np
from kmeans_pytorch import kmeans

# data
data_size, dims, num_clusters = 1000, 2, 3
x = np.random.randn(data_size, dims) / 6
x = torch.from_numpy(x)

# kmeans
cluster_ids_x, cluster_centers = kmeans(
    X=x, num_clusters=num_clusters, distance='euclidean', device=torch.device('cuda:0')

see example.ipynb for a more elaborate example


  • PyTorch version >= 1.0.0
  • Python version >= 3.6


install with pip:

pip install kmeans-pytorch

Installing from source

To install from source and develop locally:

git clone
cd kmeans_pytorch
pip install --editable .


  • useful when clustering large number of samples
  • utilizes GPU for faster matrix computations
  • support euclidean and cosine distances (for now)

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for kmeans-pytorch, version 0.2
Filename, size File type Python version Upload date Hashes
Filename, size kmeans_pytorch-0.2-py3-none-any.whl (4.2 kB) File type Wheel Python version py3 Upload date Hashes View hashes
Filename, size kmeans_pytorch-0.2.tar.gz (4.1 kB) File type Source Python version None Upload date Hashes View hashes

Supported by

Elastic Elastic Search Pingdom Pingdom Monitoring Google Google BigQuery Sentry Sentry Error logging AWS AWS Cloud computing DataDog DataDog Monitoring Fastly Fastly CDN DigiCert DigiCert EV certificate StatusPage StatusPage Status page