Implementation:Kornia Kornia KMeans

Knowledge Sources	Kornia
Domains	Vision, Clustering, Machine_Learning
Last Updated	2026-02-09 15:00 GMT

Overview

Implements the K-Means clustering algorithm with Euclidean distance on PyTorch tensors, supporting GPU execution.

Description

The KMeans class in the Kornia contrib package implements the K-Means clustering algorithm using Euclidean distance as the similarity measure. It supports initialization from random data samples or from user-provided cluster centers. The fit method iteratively assigns data points to their nearest cluster center and recomputes centers until either the shift in centers falls below a tolerance threshold or a maximum number of iterations is reached. The predict method assigns new data points to the closest learned cluster center. All computations use PyTorch tensors and can run on GPU. The implementation handles edge cases such as empty clusters by reassigning them to random data points.

Usage

Import this module when you need to perform K-Means clustering on tensor data, for example for color quantization, feature space partitioning, or unsupervised grouping of image patches or descriptors.

Code Reference

Source Location

• Repository: Kornia
• File: kornia/contrib/kmeans.py
• Lines: 1-198

Signature

class KMeans:
    def __init__(
        self,
        num_clusters: int,
        cluster_centers: torch.Tensor | None,
        tolerance: float = 10e-4,
        max_iterations: int = 0,
        seed: int | None = None,
    ) -> None: ...

    @property
    def cluster_centers(self) -> torch.Tensor: ...

    @property
    def cluster_assignments(self) -> torch.Tensor: ...

    def fit(self, X: torch.Tensor) -> None: ...

    def predict(self, x: torch.Tensor) -> torch.Tensor: ...

Import

from kornia.contrib import KMeans

I/O Contract

Inputs (__init__)

Name	Type	Required	Description
num_clusters	int	Yes	Number of clusters to assign the data to (must be > 0)
cluster_centers	torch.Tensor or None	Yes	Optional starting cluster centers with shape (C, D); if None, centers are initialized randomly from data
tolerance	float	No	Convergence threshold; algorithm stops when center shift squared is below this value (default: 10e-4)
max_iterations	int	No	Maximum number of iterations; 0 means no iteration limit (default: 0)
seed	int or None	No	Random seed for reproducibility

Inputs (fit)

Name	Type	Required	Description
X	torch.Tensor	Yes	2D input tensor to cluster with shape (N, D), where N is the number of samples and D is the feature dimension

Inputs (predict)

Name	Type	Required	Description
x	torch.Tensor	Yes	2D tensor of new data points with shape (M, D) to assign to clusters

Outputs

Name	Type	Description
cluster_centers (property)	torch.Tensor	Final cluster centers with shape (C, D) after fitting
cluster_assignments (property)	torch.Tensor	1D tensor of cluster IDs assigned to each training sample after fitting
predictions (from predict)	torch.Tensor	1D tensor of cluster IDs assigned to each input data point

Usage Examples

import torch
from kornia.contrib import KMeans

# Initialize KMeans with 3 clusters
kmeans = KMeans(num_clusters=3, cluster_centers=None, tolerance=10e-4, max_iterations=100, seed=42)

# Fit on training data
data = torch.rand(1000, 5)  # 1000 samples, 5 features
kmeans.fit(data)

# Access learned cluster centers
print(kmeans.cluster_centers.shape)      # torch.Size([3, 5])
print(kmeans.cluster_assignments.shape)  # torch.Size([1000])

# Predict cluster assignments for new data
new_data = torch.rand(10, 5)
predictions = kmeans.predict(new_data)
print(predictions)  # tensor of cluster IDs, shape (10,)

# Initialize with known cluster centers
initial_centers = torch.tensor([[0.0, 0.0], [1.0, 1.0], [2.0, 2.0]])
kmeans2 = KMeans(num_clusters=3, cluster_centers=initial_centers)
kmeans2.fit(torch.rand(500, 2))

Related Pages

• Principle:Kornia_Kornia_KMeans_Clustering