Install & Compatibility
Where this runs
tested against v0.3 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
muslpy 3.10–3.910 runs
installs and imports cleanly · install 0.0s · import 0.000s · 17.8MB
glibcpy 3.10–3.910 runs
installs and imports cleanly · install 1.6s · import 0.000s · 18MB
16MB installed
● package 16MB
Code
Verified usage
Verified import paths — ran on the pinned version, not inferred.
kmeans
✓ from kmeans_pytorch import kmeans
✗ from kmeans_pytorch import kmeans
This example demonstrates how to generate sample data, run the `kmeans` algorithm, and retrieve the cluster assignments and final cluster centers. Remember to adjust the `device` parameter ('cpu' or 'cuda:0') based on your hardware.
import torch
from kmeans_pytorch import kmeans
# 0. Generate some random data
num_samples = 1000
num_features = 2
X = torch.randn(num_samples, num_features, device='cpu', dtype=torch.float)
# Add some clusters
X[:300] += 5
X[300:600] -= 5
X[600:] += torch.tensor([0, 10], dtype=torch.float)
num_clusters = 3
tolerance = 1e-4
max_iterations = 500
distance_metric = 'euclidean'
device = 'cpu' # Change to 'cuda:0' if a GPU is available
# 1. Run K-means
cluster_ids_x, cluster_centers = kmeans(
X=X,
num_clusters=num_clusters,
distance=distance_metric,
tol=tolerance,
max_iter=max_iterations,
device=device
)
print(f"Cluster IDs shape: {cluster_ids_x.shape}")
print(f"Cluster Centers shape: {cluster_centers.shape}")
print(f"First 5 cluster IDs: {cluster_ids_x[:5]}")
print(f"Cluster centers:\n{cluster_centers}")
Debug
Known issues
gotchaPerformance on large datasets will be significantly impacted if you forget to specify `device='cuda:0'` when a GPU is available. The default device is CPU, which is much slower for heavy computations.fixPass `device='cuda:0'` to the `kmeans` function to leverage GPU acceleration. Ensure your input tensor `X` is also on the correct device (e.g., `X = X.to('cuda:0')`). affects: 0.2+
gotchaMismatched data types (e.g., `torch.float` vs `torch.double`) between the input tensor `X` and internally generated tensors can cause `RuntimeError: Input type (Float) and weight type (Double) should be the same`.fixEnsure your input tensor `X` has the same `dtype` (e.g., `torch.float` or `torch.double`) as expected by PyTorch operations within the library. Explicitly cast `X` if necessary: `X = X.to(dtype=torch.float)`.
affects: 0.2+
breakingWhile not explicitly documented as breaking changes, minor releases might introduce or clarify argument names, types, and default values for the `kmeans` function. For instance, the exact default values for `tol`, `max_iter`, or the `distance` metric might subtly change.fixAlways explicitly pass all desired arguments (e.g., `num_clusters`, `distance`, `tol`, `max_iter`, `device`) rather than relying on assumed defaults to ensure consistent behavior across updates.
affects: Prior to 0.3, possibly minor revisions within 0.3.
Upgrade
Version history
0.3latest on PyPI · released Feb 3, 2020
Audit
Dependencies
torchrequiredCore dependency for tensor operations and GPU acceleration.