Built a benchmark running platform for personalized federated learning approaches

docs/configuration.md

@@ -25,7 +25,7 @@ Valid values are `true` or `false`. The default value is `false`.
 ## clients
 
 ```{admonition} **type**
-The type of the federated learning client. Valid values include `simple`, which represents a basic client who sends weight updates to the server; and `mistnet`, which is client following the MistNet algorithm.
+The type of the federated learning client. Valid values include `simple`, which represents a basic client who sends weight updates to the server; and `mistnet`, which is client following the MistNet algorithm; and `simple_personalized`, which controls the personalized learning logic.
 ```
 
 ```{admonition} **total_clients**
@@ -136,9 +136,12 @@ A list of processors for the client to apply on the payload before receiving it
 
 ```
 
-
 ```{admonition} participant_clients_ratio
-Percentage of clients participating in federated training out of all clients.
+Percentage of clients participating in federated training out of all clients. The value should range from 0 to 1.
+```
+
+```{admonition} persist_personalized_model
+Whether the status of the trained personalized model will be persisted for subsequent training. Setting `True` will force each client to load the previously trained personalized model before any personalization learning. 
 
 ```
 
@@ -274,8 +277,9 @@ The round interval for a server commanding when to perform personalization. The
 ```
 
 ```{admonition} do_personalization_group
-The group of clients that is required by the server to perform personalization. The default value is "participant", meaning the clients participating in training.
+The group of clients that is required by the server to perform personalization. There are three options, including "total", "participant", and "nonparticipant". The default value is "participant", meaning the clients participating in training will be used to perform the personalization.
 ```
+
 ## data
 
 ```{admonition} **dataset**
@@ -378,6 +382,7 @@ The number of samples in the server's test dataset when server-side evaluation i
 The type of the trainer. The following types are available:
 - `basic`: a basic trainer with a standard training loop.
 - `diff_privacy`: a trainer that supports local differential privacy in its training loop by adding noise to the gradients during each step of training.
+- `basic_personalized`: a trainer supports all-purpose personalized federated learning by adding the personalized learning process.
 
 ```{admonition} max_physical_batch_size
 The limit on the physical batch size when using the `diff_privacy` trainer.  The default value is 128. The GPU memory usage of one process training the ResNet-18 model is around 2817 MB.
@@ -539,6 +544,39 @@ For `resnet_x`, x = 18, 34, 50, 101, or 152; For `vgg_x`, x = 11, 13, 16, or 19.
 ```
 ````
 
+````{admonition} **personalized learning**
+The variables used when the client performs the personalization locally:
+
+- `personalized_epochs` 
+- `personalized_batch_size`
+- `personalized_optimizer`
+- `personalized_lr_scheduler`
+- `personalized_loss_criterion`
+- `personalized_model_name`
+- `personalized_epoch_log_interval` (Logging the training details)
+
+
+````{note}
+The definitions, configurations, and applications of these variables for personalization align with those of traditional federated learning described previously.
+````
+
+The specific hyperparameters of certain variables, including `model`, `personalized_optimizer`, `personalized_lr_scheduler`, and `personalized_loss_criterion`, can be configured in the parameters section of the configuration file, as demonstrated below:
+
+```
+parameters:
+    model:
+        num_classes: 10
+    personalized_optimizer:
+        lr: 0.1
+        momentum: 0.9
+        weight_decay: 0.0
+```
+
+
+````
+
+
+
 ## algorithm
 
 ```{admonition} **type**

examples/fedavg_finetune/fedavg_finetune.py

@@ -0,0 +1,31 @@
+"""
+An implementation of the personalized learning variant of FedAvg.
+
+Such an variant of FedAvg is recently mentioned and discussed in work [1].
+
+[1] Liam Collins, et al., "Exploiting shared representations for personalized federated learning,"
+in the Proceedings of ICML 2021.
+
+    Address: https://proceedings.mlr.press/v139/collins21a.html
+
+    Code: https://github.com/lgcollins/FedRep
+
+"""
+
+from fedavg_finetune_client import Client
+from plato.servers import registry as server_registry
+
+
+def main():
+    """
+    A Plato personalized federated learning training session using the FedAvg algorithm under the
+    supervised learning setting.
+    """
+    client = Client()
+    server = server_registry.get()
+
+    server.run(client)
+
+
+if __name__ == "__main__":
+    main()

examples/fedavg_finetune/fedavg_finetune_MNIST_lenet5_noniid.yml

@@ -0,0 +1,122 @@
+clients:
+    # Type
+    type: simple_personalized
+
+    # The total number of clients
+    total_clients: 10
+
+    # The number of clients selected in each round
+    per_round: 3
+
+    # Should the clients compute test accuracy locally?
+    do_test: true
+
+    # the ratio of clients participanting in training
+    participant_clients_ratio: 0.5
+
+    # whether maintain and load the trained 
+    # personalied model
+    # set to be true, each client will load
+    # the latest trained personalized model
+    persist_personalized_model: False
+
+server:
+    type: fedavg_personalized
+    address: 127.0.0.1
+    port: 8000
+
+    # do personalization 
+    # for every `do_personalization_interval` rounds
+    # set -1 to do personalized only after the final round
+    do_personalization_interval: -1
+    # total, participant, nonparticipant
+    do_personalization_group: total
+
+data:
+    # The training and testing dataset
+    datasource: MNIST
+
+    # Number of samples in each partition
+    partition_size: 2000
+
+
+    # IID or non-IID?
+    sampler: noniid
+    testset_sampler: noniid
+
+    # The concentration parameter for the Dirichlet distribution
+    concentration: 0.5
+
+    # The random seed for sampling data
+    random_seed: 1
+
+trainer:
+    # The type of the trainer
+    type: basic_personalized
+
+    # The maximum number of training rounds
+    rounds: 5
+
+    # The maximum number of clients running concurrently
+    max_concurrency: 2
+
+    # The target accuracy
+    target_accuracy: 0.94
+
+    # Number of epoches for local training in each communication round
+    epochs: 3
+    batch_size: 32
+    optimizer: SGD
+    lr_scheduler: MultiStepLR
+    loss_criterion: CrossEntropyLoss
+
+    # Number of epochs for local training in each communication round
+    #   The batch size for the downstream task can be larger as it
+    # does not utilize the
+    personalized_epochs: 10
+    personalized_batch_size: 12
+    personalized_optimizer: SGD
+    personalized_lr_scheduler: StepLR
+    personalized_loss_criterion: CrossEntropyLoss
+
+    pers_epoch_log_interval: 2
+
+    # The machine learning model
+    model_name: lenet5
+    personalized_model_name: lenet5
+
+algorithm:
+    # Aggregation algorithm
+    type: fedavg
+
+parameters:
+    optimizer:
+        lr: 0.05
+        momentum: 0.9
+        weight_decay: 0.000001
+
+    personalized_optimizer:
+        lr: 0.1
+        momentum: 0.9
+        weight_decay: 0.0
+
+    learning_rate:
+        gamma: 0.1
+        milestone_steps: 150ep,350ep
+
+    personalized_learning_rate:
+        gamma: 0.3
+        step_size: 10
+
+    loss_criterion:
+        label_smoothing: 0.5
+
+    personalized_loss_criterion:
+        label_smoothing: 0.5
+
+    model:
+        num_classes: 10
+
+    personalized_model:
+        num_classes: 10
+

examples/fedavg_finetune/fedavg_finetune_client.py

@@ -0,0 +1,35 @@
+"""
+A personalized federated learning client using FedAvg.
+
+"""
+import logging
+
+
+from plato.clients import simple_personalized
+from plato.config import Config
+from plato.utils import fonts
+
+
+class Client(simple_personalized.Client):
+    """A personalized federated learning trainer using the FedAvg algorithm."""
+
+    def load_personalized_model(self) -> None:
+        """Load the personalized model.
+
+        Each client of FedAvg will directly utilize the recevied global model as the
+        personalized model.
+        """
+        logging.info(
+            fonts.colourize(
+                "[Client #%d] assings the received model [%s] to personalized model [%s].",
+                colour="blue",
+            ),
+            self.client_id,
+            Config().trainer.model_name,
+            Config().trainer.personalized_model_name,
+        )
+
+        # load the received model to be personalized model
+        self.trainer.personalized_model.load_state_dict(
+            self.trainer.model.state_dict(), strict=True
+        )

plato/clients/registry.py

@@ -10,10 +10,12 @@
 from plato.clients import (
     simple,
     mistnet,
+    simple_personalized
 )
 
 registered_clients = {
     "simple": simple.Client,
+    "simple_personalized": simple_personalized.Client,
     "mistnet": mistnet.Client,
 }