Sync dev

.github/workflows/docker-publish.yml

@@ -15,6 +15,7 @@ env:
   # Use docker.io for Docker Hub if empty
   GHCR_REGISTRY: ghcr.io
   ALIYUN_REGISTRY: registry.cn-hangzhou.aliyuncs.com
+  DOCKERHUB_REGISTRY: docker.io
   # github.repository as <account>/<repo>
   IMAGE_NAME: ${{ github.repository }}
 
@@ -99,4 +100,34 @@ jobs:
           tags: ${{ steps.meta-aliyun.outputs.tags }}
           labels: ${{ steps.meta-aliyun.outputs.labels }}
           cache-from: type=gha
+          cache-to: type=gha,mode=max
+
+      # Login against a Docker registry except on PR
+      # https://github.com/docker/login-action
+      - name: Log into registry ${{ env.DOCKERHUB_REGISTRY }}
+        if: github.event_name != 'pull_request'
+        uses: docker/login-action@343f7c4344506bcbf9b4de18042ae17996df046d # v3.0.0
+        with:
+          registry: ${{ env.DOCKERHUB_REGISTRY }}
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+      # Extract metadata (tags, labels) for Docker
+      # https://github.com/docker/metadata-action
+      - name: Extract Docker metadata
+        id: meta-dockerhub
+        uses: docker/metadata-action@96383f45573cb7f253c731d3b3ab81c87ef81934 # v5.0.0
+        with:
+          images: ${{ env.DOCKERHUB_REGISTRY }}/${{ env.IMAGE_NAME }}
+
+      # Build and push Docker image with Buildx (don't push on PR)
+      # https://github.com/docker/build-push-action
+      - name: Build and push Docker image
+        id: build-and-push-to-Dockerhub-container-registry
+        uses: docker/build-push-action@0565240e2d4ab88bba5387d719585280857ece09 # v5.0.0
+        with:
+          context: .
+          push: ${{ github.event_name != 'pull_request' }}
+          tags: ${{ steps.meta-Dockerhub.outputs.tags }}
+          labels: ${{ steps.meta-Dockerhub.outputs.labels }}
+          cache-from: type=gha
           cache-to: type=gha,mode=max

CITATION.cff

@@ -0,0 +1,25 @@
+# This CITATION.cff file was generated with cffinit.
+# Visit https://bit.ly/cffinit to generate yours today!
+
+cff-version: 1.2.0
+title: 'FL-bench: A federated learning benchmark for solving image classification tasks'
+message: >-
+  If you use this software, please cite it using the
+  metadata from this file.
+type: software
+authors:
+  - given-names: Jiahao
+    family-names: Tan
+    email: karhoutam@qq.com
+    affiliation: Shenzhen University
+  - given-names: Xinpeng
+    family-names: Wang
+    affiliation: 'The Chinese University of Hong Kong, Shenzhen'
+    email: 223015056@link.cuhk.edu.cn
+repository-code: 'https://github.com/KarhouTam/FL-bench'
+abstract: >-
+  Benchmark of federated learning that aim solving image
+  classification tasks.
+keywords:
+  - federated learning
+license: GPL-2.0

README.md

@@ -92,6 +92,8 @@
 
 - ***MetaFed*** -- [MetaFed: Federated Learning among Federations with Cyclic Knowledge Distillation for Personalized Healthcare](http://arxiv.org/abs/2206.08516) (IJCAI'22)
 
+- ***FedRoD*** -- [On Bridging Generic and Personalized Federated Learning for Image Classification](https://arxiv.org/abs/2107.00778) (ICLR'22)
+
 ### FL Domain Generalization Methods
 
 - ***FedSR*** -- [FedSR: A Simple and Effective Domain Generalization Method for Federated Learning](https://openreview.net/forum?id=mrt90D00aQX) (NIPS'22)
@@ -251,3 +253,24 @@ Medical Image Datasets
 - [*COVID-19*](https://www.researchgate.net/publication/344295900_Curated_Dataset_for_COVID-19_Posterior-Anterior_Chest_Radiography_Images_X-Rays) (3 x 244 x 224, 4 classes)
 
 - [*Organ-S/A/CMNIST*](https://medmnist.com/) (1 x 28 x 28, 11 classes)
+
+## Citation 🧐
+
+```
+@software{Tan_FL-bench,
+          author = {Tan, Jiahao and Wang, Xinpeng},
+          license = {GPL-2.0},
+          title = {{FL-bench: A federated learning benchmark for solving image classification tasks}},
+          url = {https://github.com/KarhouTam/FL-bench}
+}
+
+@misc{tan2023pfedsim,
+      title={pFedSim: Similarity-Aware Model Aggregation Towards Personalized Federated Learning}, 
+      author={Jiahao Tan and Yipeng Zhou and Gang Liu and Jessie Hui Wang and Shui Yu},
+      year={2023},
+      eprint={2305.15706},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG}
+}
+
+```

src/client/apfl.py

@@ -69,7 +69,6 @@ def fit(self):
                 logit_g = self.model(x)
                 logit_p = self.alpha * logit_l + (1 - self.alpha) * logit_g.detach()
                 loss = self.criterion(logit_p, y)
-                self.optimizer.zero_grad()
                 loss.backward()
                 self.optimizer.step()
 
@@ -93,9 +92,9 @@ def update_alpha(self):
         self.alpha.data -= self.args.local_lr * alpha_grad
         self.alpha.clip_(0, 1.0)
 
-    def evaluate(self, model=None, test_flag=False):
+    def evaluate(self):
         return super().evaluate(
-            MixedModel(self.local_model, self.model, alpha=self.alpha), test_flag
+            model=MixedModel(self.local_model, self.model, alpha=self.alpha)
         )
 
 

src/client/ditto.py

@@ -60,5 +60,5 @@ def fit(self):
                     )
                 self.optimizer.step()
 
-    def evaluate(self, model=None, test_flag=False) -> Dict[str, float]:
-        return super().evaluate(self.pers_model, test_flag)
+    def evaluate(self):
+        return super().evaluate(self.pers_model)

src/client/fedavg.py

@@ -12,6 +12,7 @@
 PROJECT_DIR = Path(__file__).parent.parent.parent.absolute()
 
 from src.utils.tools import trainable_params, evalutate_model, Logger
+from src.utils.metrics import Metrics
 from src.utils.models import DecoupledModel
 from src.utils.constants import DATA_MEAN, DATA_STD
 from data.utils.datasets import DATASETS
@@ -69,6 +70,7 @@ def __init__(
         self.trainset: Subset = Subset(self.dataset, indices=[])
         self.valset: Subset = Subset(self.dataset, indices=[])
         self.testset: Subset = Subset(self.dataset, indices=[])
+        self.test_flag = False
 
         self.model = model.to(self.device)
         self.local_epoch = self.args.local_epoch
@@ -115,50 +117,35 @@ def train_and_log(self, verbose=False) -> Dict[str, Dict[str, float]]:
         Returns:
             Dict[str, Dict[str, float]]: The logging info, which contains metric stats.
         """
-        eval_results = {
-            "before": {
-                "train": {"loss": 0, "correct": 0, "size": 0},
-                "val": {"loss": 0, "correct": 0, "size": 0},
-                "test": {"loss": 0, "correct": 0, "size": 0},
-            },
-            "after": {
-                "train": {"loss": 0, "correct": 0, "size": 0},
-                "val": {"loss": 0, "correct": 0, "size": 0},
-                "test": {"loss": 0, "correct": 0, "size": 0},
-            },
+        eval_metrics = {
+            "before": {"train": Metrics(), "val": Metrics(), "test": Metrics()},
+            "after": {"train": Metrics(), "val": Metrics(), "test": Metrics()},
         }
-        eval_results["before"] = self.evaluate()
+        eval_metrics["before"] = self.evaluate()
         if self.local_epoch > 0:
             self.fit()
             self.save_state()
-            eval_results["after"] = self.evaluate()
+            eval_metrics["after"] = self.evaluate()
         if verbose:
-            colors = {"train": "yellow", "val": "green", "test": "cyan"}
-            for split, flag, subset in [
-                ["train", self.args.eval_train, self.trainset],
-                ["val", self.args.eval_val, self.valset],
-                ["test", self.args.eval_test, self.testset],
+            for split, color, flag, subset in [
+                ["train", "yellow", self.args.eval_train, self.trainset],
+                ["val", "green", self.args.eval_val, self.valset],
+                ["test", "cyan", self.args.eval_test, self.testset],
             ]:
                 if len(subset) > 0 and flag:
                     self.logger.log(
                         "client [{}] [{}]({})  loss: {:.4f} -> {:.4f}   accuracy: {:.2f}% -> {:.2f}%".format(
                             self.client_id,
-                            colors[split],
+                            color,
                             split,
-                            eval_results["before"][split]["loss"]
-                            / eval_results["before"][split]["size"],
-                            eval_results["after"][split]["loss"]
-                            / eval_results["after"][split]["size"],
-                            eval_results["before"][split]["correct"]
-                            / eval_results["before"][split]["size"]
-                            * 100.0,
-                            eval_results["after"][split]["correct"]
-                            / eval_results["after"][split]["size"]
-                            * 100.0,
+                            eval_metrics["before"][split].loss,
+                            eval_metrics["after"][split].loss,
+                            eval_metrics["before"][split].accuracy,
+                            eval_metrics["after"][split].accuracy,
                         )
                     )
 
-        return eval_results
+        return eval_metrics
 
     def set_parameters(self, new_parameters: OrderedDict[str, torch.Tensor]):
         """Load model parameters received from the server.
@@ -256,104 +243,81 @@ def fit(self):
                 self.optimizer.step()
 
     @torch.no_grad()
-    def evaluate(
-        self, model: torch.nn.Module = None, force_eval=False
-    ) -> Dict[str, float]:
+    def evaluate(self, model: torch.nn.Module = None) -> Dict[str, Metrics]:
         """The evaluation function. Would be activated before and after local training if `eval_test = True` or `eval_train = True`.
 
         Args:
             model (torch.nn.Module, optional): The target model needed evaluation (set to `None` for using `self.model`). Defaults to None.
             force_eval (bool, optional): Set as `True` when the server asking client to evaluate model.
         Returns:
-            Dict[str, float]: The evaluation metric stats.
+            Dict[str, Metrics]: The evaluation metric stats.
         """
         # disable train data transform while evaluating
         self.dataset.enable_train_transform = False
 
         target_model = self.model if model is None else model
         target_model.eval()
-        train_loss, val_loss, test_loss = 0, 0, 0
-        train_correct, val_correct, test_correct = 0, 0, 0
-        train_size, val_size, test_size = 0, 0, 0
+        train_metrics = Metrics()
+        val_metrics = Metrics()
+        test_metrics = Metrics()
         criterion = torch.nn.CrossEntropyLoss(reduction="sum")
 
         if len(self.testset) > 0 and self.args.eval_test:
-            test_loss, test_correct, test_size = evalutate_model(
+            test_metrics = evalutate_model(
                 model=target_model,
                 dataloader=self.testloader,
                 criterion=criterion,
                 device=self.device,
             )
 
-        if len(self.valset) > 0 and (force_eval or self.args.eval_val):
-            val_loss, val_correct, val_size = evalutate_model(
+        if len(self.valset) > 0 and self.args.eval_val:
+            val_metrics = evalutate_model(
                 model=target_model,
                 dataloader=self.valloader,
                 criterion=criterion,
                 device=self.device,
             )
 
         if len(self.trainset) > 0 and self.args.eval_train:
-            train_loss, train_correct, train_size = evalutate_model(
+            train_metrics = evalutate_model(
                 model=target_model,
                 dataloader=self.trainloader,
                 criterion=criterion,
                 device=self.device,
             )
 
         self.dataset.enable_train_transform = True
-        return {
-            "train": {
-                "loss": train_loss,
-                "correct": train_correct,
-                "size": float(max(1, train_size)),
-            },
-            "val": {
-                "loss": val_loss,
-                "correct": val_correct,
-                "size": float(max(1, val_size)),
-            },
-            "test": {
-                "loss": test_loss,
-                "correct": test_correct,
-                "size": float(max(1, test_size)),
-            },
-        }
+        return {"train": train_metrics, "val": val_metrics, "test": test_metrics}
 
     def test(
         self, client_id: int, new_parameters: OrderedDict[str, torch.Tensor]
-    ) -> Dict[str, Dict[str, float]]:
+    ) -> Dict[str, Dict[str, Metrics]]:
         """Test function. Only be activated while in FL test round.
 
         Args:
             client_id (int): The ID of client.
             new_parameters (OrderedDict[str, torch.Tensor]): The FL model parameters.
 
         Returns:
-            Dict[str, Dict[str, float]]: the evalutaion metrics stats.
+            Dict[str, Dict[str, Metrics]]: the evalutaion metrics stats.
         """
+        self.test_flag = True
         self.client_id = client_id
         self.load_dataset()
         self.set_parameters(new_parameters)
 
-        # set `size` as 1 for avoiding NaN.
         results = {
-            "before": {
-                "train": {"loss": 0, "correct": 0, "size": 1},
-                "val": {"loss": 0, "correct": 0, "size": 1},
-                "test": {"loss": 0, "correct": 0, "size": 1},
-            },
-            "after": {
-                "train": {"loss": 0, "correct": 0, "size": 1},
-                "val": {"loss": 0, "correct": 0, "size": 1},
-                "test": {"loss": 0, "correct": 0, "size": 1},
-            },
+            "before": {"train": Metrics(), "val": Metrics(), "test": Metrics()},
+            "after": {"train": Metrics(), "val": Metrics(), "test": Metrics()},
         }
 
-        results["before"] = self.evaluate(force_eval=True)
+        results["before"] = self.evaluate()
         if self.args.finetune_epoch > 0:
+            frz_params_dict = deepcopy(self.model.state_dict())
             self.finetune()
-            results["after"] = self.evaluate(force_eval=True)
+            results["after"] = self.evaluate()
+            self.model.load_state_dict(frz_params_dict)
+        self.test_flag = False
         return results
 
     def finetune(self):

src/client/fedfomo.py

@@ -60,8 +60,7 @@ def set_parameters(self, received_params: Dict[int, List[torch.Tensor]]):
             dataloader=self.valloader,
             criterion=self.criterion,
             device=self.device,
-        )[0]
-        LOSS /= len(self.valset)
+        ).loss
         W = torch.zeros(len(received_params), device=self.device)
         self.weight_vector.zero_()
         with torch.no_grad():
@@ -75,8 +74,7 @@ def set_parameters(self, received_params: Dict[int, List[torch.Tensor]]):
                     dataloader=self.valloader,
                     criterion=self.criterion,
                     device=self.device,
-                )[0]
-                loss /= len(self.valset)
+                ).loss
                 params_diff = vectorize(params_i) - vectorized_self_params
                 w = (LOSS - loss) / (torch.norm(params_diff) + 1e-5)
                 W[i] = w