expose aggregate_fn to users for overwriting (#2539)

NVIDIA · Apr 30, 2024 · 593a11d · 593a11d
1 parent ce79609
commit 593a11d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/nvflare/app_common/workflows/base_fedavg.py b/nvflare/app_common/workflows/base_fedavg.py
@@ -107,7 +107,7 @@ def _check_results(results: List[FLModel]):
             raise ValueError(f"Result from client(s) {empty_clients} is empty!")
 
     @staticmethod
-    def _aggregate_fn(results: List[FLModel]) -> FLModel:
+    def aggregate_fn(results: List[FLModel]) -> FLModel:
         aggregation_helper = WeightedAggregationHelper()
         for _result in results:
             aggregation_helper.add(
@@ -141,7 +141,7 @@ def aggregate(self, results: List[FLModel], aggregate_fn=None) -> FLModel:
         self._check_results(results)
 
         if not aggregate_fn:
-            aggregate_fn = self._aggregate_fn
+            aggregate_fn = self.aggregate_fn
 
         self.info(f"aggregating {len(results)} update(s) at round {self.current_round}")
         try:

diff --git a/nvflare/app_common/workflows/fedavg.py b/nvflare/app_common/workflows/fedavg.py
@@ -57,8 +57,8 @@ def run(self) -> None:
             results = self.send_model_and_wait(targets=clients, data=model)
 
             aggregate_results = self.aggregate(
-                results, aggregate_fn=None
-            )  # if no `aggregate_fn` provided, default `WeightedAggregationHelper` is used
+                results, aggregate_fn=self.aggregate_fn
+            )  # using default aggregate_fn with `WeightedAggregationHelper`. Can overwrite self.aggregrate_fn with signature Callable[List[FLModel], FLModel]
 
             model = self.update_model(model, aggregate_results)