KMeans-GPU: A PyTorch Module for KMeans.
Project description
kmeans-gpu
kmeans-gpu with pytorch (batch version). It is faster than sklearn.cluster.KMeans.
You can easily use KMeans
as a nn.Module
, and embed into your network structure.
Install
- From Git:
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__)"
- From PyPI:
pip install kmeans-gpu
# check installation
python -c "import kmeans_gpu; print(kmeans_gpu.__version__)"
Demo
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])
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
kmeans_gpu-0.0.1.tar.gz
(4.7 kB
view hashes)
Built Distributions
kmeans_gpu-0.0.1-py3.7.egg
(8.2 kB
view hashes)
Close
Hashes for kmeans_gpu-0.0.1-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 3a771da60f968626084023ad15d7b2a04ff17e3a9c3e1a25accbe527905c0d1f |
|
MD5 | 4e52b0336bfdf3f3a8c9b2600bfe7d87 |
|
BLAKE2b-256 | 846824c9c23c41ee1509a8e83ba1caa02329407cbbed5e3954ffa2010fc49733 |