From ebdb2ed451c33a37253705928be40a731657f145 Mon Sep 17 00:00:00 2001
From: futabato <01futabato10@gmail.com>
Date: Fri, 26 Jul 2024 17:57:31 +0900
Subject: [PATCH] :sparkles: Introduce krum aggregation

---
 src/federatedlearning/main.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/src/federatedlearning/main.py b/src/federatedlearning/main.py
index af3b918..643d19b 100644
--- a/src/federatedlearning/main.py
+++ b/src/federatedlearning/main.py
@@ -30,6 +30,7 @@
 )
 from federatedlearning.server.aggregations.aggregators import (
     average_weights,
+    krum,
     median_weights,
 )
 from federatedlearning.server.inferencing import inference
@@ -283,6 +284,10 @@ def main(cfg: DictConfig) -> float:  # noqa: C901
                 global_weights = average_weights(local_weights)
             elif cfg.federatedlearning.aggregation == "median":
                 global_weights = median_weights(local_weights)
+            elif cfg.federatedlearning.aggregation == "krum":
+                global_weights = krum(
+                    local_weights, cfg.federatedlearning.num_byzantines
+                )
             else:
                 global_weights = average_weights(local_weights)
             # Save updated global model weights