pyagc.clusters.TritonKMeans

class TritonKMeans(*args, **kwargs)[source]

Bases: object

Fast K-Means clustering implemented with Triton GPU kernels.

This implementation provides an interface compatible with TorchKMeans while leveraging Triton kernels for improved performance.

Parameters:

  • metric (str, default='euclidean') – Distance metric to use. Options: ‘euclidean’, ‘cosine’, ‘dot’

  • init (str or torch.Tensor, default='k-means++') – Method for initialization: ‘k-means++’, ‘random’ or user-specified tensor of shape (n_clusters, n_features).

  • random_state (int, optional) – Random seed for centroid initialization.

  • n_clusters (int, default=8) – Number of clusters (k).

  • n_init (int, default=10) – Number of times the algorithm will be run with different centroid seeds. The final result will be the best output of n_init consecutive runs.

  • max_iter (int, default=300) – Maximum number of iterations for a single run.

  • tol (float, default=1e-4) – Relative tolerance with regards to inertia to declare convergence.

  • verbose (bool, default=False) – Whether to print per-iteration info.

  • dtype (torch.dtype, optional) – Compute data type for algorithm.

  • device (torch.device, optional) – Target device. Defaults to “cuda:0” when available. Currently, only CUDA devices are supported.

  • distributed (bool, default=False) – Reserved for future distributed training support (currently not implemented).

__init__(*args, **kwargs)[source]

Methods

__init__(*args, **kwargs)

fit(X)

Fit k-means clustering.

fit_predict(X)

Fit and predict cluster labels.

fit_transform(X)

Fit and transform.

predict(X[, soft])

Predict cluster labels.

transform(X)

Transform to cluster-distance space.
fit(X)[source]

Fit k-means clustering.

fit_predict(X)[source]

Fit and predict cluster labels.

predict(X, soft=False)[source]

Predict cluster labels.

transform(X)[source]

Transform to cluster-distance space.

fit_transform(X)[source]

Fit and transform.