tff.learning.algorithms.build_fed_kmeans  |  TensorFlow Federated (original) (raw)

Builds a learning process for federated k-means clustering.

tff.learning.algorithms.build_fed_kmeans(
    num_clusters: int,
    data_shape: tuple[int, ...],
    random_seed: Optional[tuple[int, int]] = None,
    distributor: Optional[tff.learning.templates.DistributionProcess] = None,
    sum_aggregator: Optional[tff.aggregators.UnweightedAggregationFactory] = None
) -> tff.learning.templates.LearningProcess

This function creates a tff.learning.templates.LearningProcess that performs federated k-means clustering. Specifically, this performs mini-batch k-means clustering. Note that mini-batch k-means only processes a mini-batch of the data at each round, and updates clusters in a weighted manner based on how many points in the mini-batch were assigned to each cluster. In the federated version, clients do the assignment of each of their point locally, and the server updates the clusters. Conceptually, the "mini-batch" being used is the union of all client datasets involved in a given round.

The learning process has the following methods inherited fromtff.learning.templates.LearningProcess:

• initialize: A tff.Computation with the functional type signature( -> S@SERVER), where S is a LearningAlgorithmState representing the initial state of the server.
• next: A tff.Computation with the functional type signature(<S@SERVER, {B*}@CLIENTS> -> <L@SERVER>) where S is aLearningAlgorithmState whose type matches the output of initializeand {B*}@CLIENTS represents the client datasets. The output L is atff.learning.templates.LearningProcessOutput containing the state Sand metrics computed during training.
• get_model_weights: A tff.Computation with type signature (S -> W), where W represents the current k-means centroids.
• set_model_weights: A tff.Computation with type signature(<S, M> -> S), where S is atff.learning.templates.LearningAlgorithmState whose type matches the output of initialize and M a new set of k-means centroids.

Here, S is a tff.learning.templates.LearningAlgorithmState. The centroidsW is a tensor representing the current centroids, and is of shape(num_clusters,) + data_shape. The datasets {B*} must have elements of shape data_shape, and not employ batching.

The centroids are updated at each round by assigning all clients' points to the nearest centroid, and then summing these points according to these centroids. The centroids are then updated at the server based on these points. To do so, we keep track of how many points have been assigned to each centroid overall, as a float tensor of shape (num_clusters,). This information can be found in state.finalizer. Note that we begin with a "pseudo-count" of 1, in order to ensure that the centroids do not collapse to zero.