KMeansModel.forward

KMeansModel.forward(data)[source]

Computes the distance of each datapoint to each centroid as well as the "inertia", the squared distance of each datapoint to its closest centroid.

Parameters

data (Tensor) -- A tensor of shape [num_datapoints, num_features] for which to compute the distances and inertia.

Return type

Tuple[Tensor, Tensor, Tensor]

Returns

  • A tensor of shape [num_datapoints, num_centroids] with the distance from each datapoint to each centroid.

  • A tensor of shape [num_datapoints] with the assignments, i.e. the indices of each datapoint's closest centroid.

  • A tensor of shape [num_datapoints] with the inertia (squared distance to the closest centroid) of each datapoint.