expose aggregate_fn to users for overwriting

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

Hide file tree

Showing 2 changed files with 3 additions and 3 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(

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)