https://github.com/densechen/kmeans-gpu
kmeans-gpu with pytorch support
https://github.com/densechen/kmeans-gpu
Last synced: 5 months ago
JSON representation
kmeans-gpu with pytorch support
- Host: GitHub
- URL: https://github.com/densechen/kmeans-gpu
- Owner: densechen
- License: mit
- Created: 2021-07-03T00:26:01.000Z (almost 5 years ago)
- Default Branch: main
- Last Pushed: 2021-07-08T23:29:01.000Z (almost 5 years ago)
- Last Synced: 2025-09-28T09:06:13.334Z (9 months ago)
- Language: Jupyter Notebook
- Size: 105 KB
- Stars: 13
- Watchers: 1
- Forks: 3
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# kmeans-gpu
kmeans-gpu with pytorch (batch version). It is faster than sklearn.cluster.KMeans.
What's more, it is a differential operation which will back-propagate gradient to previous layers.
You can easily use `KMeans` as a `nn.Module`, and embed into your network structure.
## Install
1. From Git:
```bash
git clone git@github.com:densechen/kmeans-gpu.git
cd kmeans-gpu
pip install -r requirements.txt
python setup.py install
# check installation
python -c "import kmeans_gpu; print(kmeans_gpu.__version__)"
```
2. From PyPI:
```bash
pip install kmeans-gpu
# check installation
python -c "import kmeans_gpu; print(kmeans_gpu.__version__)"
```
## Demo
```python
from kmeans_gpu import KMeans
import torch
# Config
batch_size = 128
feature_dim = 1024
pts_dim = 3
num_pts = 256
num_cluster = 15
# Create data
features = torch.randn(batch_size, feature_dim, num_pts)
# Pay attention to the different dimension order between features and points.
points = torch.randn(batch_size, num_pts, pts_dim)
# Create KMeans Module
kmeans = KMeans(
n_clusters=num_cluster,
max_iter=100,
tolerance=1e-4,
distance='euclidean',
sub_sampling=None,
max_neighbors=15,
)
# Forward
centroids, features = kmeans(points, features)
print(centroids.shape, features.shape)
# output:
# >>> torch.Size([128, 15, 3]) torch.Size([128, 1024, 15])
```