torch_kmeans.clustering.constr_kmeans module

class torch_kmeans.clustering.constr_kmeans.ConstrainedKMeans(init_method: str = 'rnd', num_init: int = 8, max_iter: int = 100, distance: ~torch_kmeans.utils.distances.BaseDistance = <class 'torch_kmeans.utils.distances.LpDistance'>, p_norm: int = 2, tol: float = 0.0001, n_clusters: ~typing.Optional[int] = 8, verbose: bool = True, seed: ~typing.Optional[int] = 123, n_priority_trials_before_fall_back: int = 5, raise_infeasible: bool = True, **kwargs)[source]

Bases: KMeans

Implements constrained k-means clustering. Priority implementation is based on the method of

Paper:

Geetha, S., G. Poonthalir, and P. T. Vanathi. “Improved k-means algorithm for capacitated clustering problem.” INFOCOMP Journal of Computer Science 8.4 (2009)

Parameters

  • init_method (str) – Method to initialize cluster centers: [‘rnd’, ‘topk’, ‘k-means++’, ‘ckm++’] (default: ‘rnd’)

  • num_init (int) – Number of different initial starting configurations, i.e. different sets of initial centers (default: 8).

  • max_iter (int) – Maximum number of iterations (default: 100).

  • distance (BaseDistance) – batched distance evaluator (default: LpDistance).

  • p_norm (int) – norm for lp distance (default: 2).

  • tol (float) – Relative tolerance with regards to Frobenius norm of the difference in the cluster centers of two consecutive iterations to declare convergence. (default: 1e-4)

  • n_clusters (Optional[int]) – Default number of clusters to use if not provided in call (optional, default: 8).

  • verbose (bool) – Verbosity flag to print additional info (default: True).

  • seed (Optional[int]) – Seed to fix random state for randomized center inits (default: 123).

  • n_priority_trials_before_fall_back (int) – Number of trials trying to assign samples to constrained clusters based on priority values before falling back to assigning the node with the highest weight to a cluster which can still accommodate it or the dummy cluster otherwise. (default: 5)

  • raise_infeasible (bool) – if set to False, will only display a warning instead of raising an error (default: True)

  • **kwargs – additional key word arguments for the distance function.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

INIT_METHODS = ['rnd', 'k-means++', 'topk', 'ckm++']
NORM_METHODS = []
predict(x: Tensor, weights: Tensor, **kwargs) → LongTensor[source]

Predict the closest cluster each sample in X belongs to.

Parameters

  • x (Tensor) – input features/coordinates (BS, N, D)

  • weights (Tensor) – normalized weight for each sample (BS, N)

  • **kwargs – additional kwargs for assignment procedure

Returns

batch tensor of cluster labels for each sample (BS, N)

Return type

LongTensor

training: bool