Fixed issue with device during evaluation

- Fixed model in kafé configuration
- Added unit test for kafé

.gitignore

@@ -14,4 +14,5 @@ fluke_fl.egg-info/
 tests/*.ipynb
 runs/
 repr_results/
-checkpoint*/
+checkpoint*/
+chk_*

configs/kafe.yaml

@@ -13,5 +13,5 @@ hyperparameters:
   server:
     weighted: true
     bandwidth: 1.0
-  model: CNN_Mnist
+  model: MNIST_2NN
 name: fluke.algorithms.kafe.Kafe

fluke/algorithms/fedhp.py

@@ -123,7 +123,7 @@ def fit(self, override_local_epochs: int = 0) -> float:
     def evaluate(self, evaluator: Evaluator, test_set: FastDataLoader) -> dict[str, float]:
         if test_set is not None and self.initial_prototypes is not None:
             model = FedHPModel(self.model)
-            return evaluator.evaluate(self._last_round, model, test_set)
+            return evaluator.evaluate(self._last_round, model, test_set, device=self.device)
         return {}
 
     def finalize(self) -> None:

fluke/algorithms/fednh.py

@@ -138,7 +138,7 @@ def fit(self, override_local_epochs: int = 0) -> float:
     def evaluate(self, evaluator: Evaluator, test_set: FastDataLoader) -> dict[str, float]:
         if test_set is not None and self.model is not None:
             model = ArgMaxModule(self.model)
-            return evaluator.evaluate(self._last_round, model, test_set)
+            return evaluator.evaluate(self._last_round, model, test_set, device=self.device)
         return {}
 
     def finalize(self) -> None:
@@ -225,7 +225,7 @@ def aggregate(self, eligible: Iterable[PFLClient]) -> None:
     def evaluate(self, evaluator: Evaluator, test_set: FastDataLoader) -> dict[str, float]:
         if self.test_set is not None:
             model = ArgMaxModule(self.model)
-            return evaluator.evaluate(self.rounds + 1, model, self.test_set)
+            return evaluator.evaluate(self.rounds + 1, model, self.test_set, device=self.device)
         return {}
 
 

fluke/algorithms/fedproto.py

@@ -151,7 +151,7 @@ def fit(self, override_local_epochs: int = 0) -> float:
     def evaluate(self, evaluator: Evaluator, test_set: FastDataLoader) -> dict[str, float]:
         if test_set is not None and self.prototypes[0] is not None:
             model = FedProtoModel(self.model, self.prototypes, self.device)
-            return evaluator.evaluate(self._last_round, model, test_set)
+            return evaluator.evaluate(self._last_round, model, test_set, device=self.device)
         return {}
 
     def finalize(self) -> None:

fluke/algorithms/fedrod.py

@@ -138,7 +138,8 @@ def evaluate(self, evaluator: Evaluator, test_set: FastDataLoader) -> dict[str,
         if test_set is not None and self.model is not None and self.inner_model is not None:
             return evaluator.evaluate(self._last_round,
                                       RODModel(self.model, self.inner_model),
-                                      test_set)
+                                      test_set,
+                                      device=self.device)
         return {}
 
 

fluke/client.py

@@ -254,7 +254,7 @@ def evaluate(self,
             the results.
         """
         if test_set is not None and self.model is not None:
-            return evaluator.evaluate(self._last_round, self.model, test_set)
+            return evaluator.evaluate(self._last_round, self.model, test_set, device=self.device)
         return {}
 
     def finalize(self) -> None:
@@ -319,7 +319,7 @@ def __str__(self) -> str:
         hpstr = ", ".join([f"{h}={str(v)}" for h, v in self.hyper_params.items()])
         hpstr = ", " + hpstr if hpstr else ""
         return f"{self.__class__.__name__}[{self._index}](optim={self.optimizer_cfg}, " + \
-               f"batch_size={self.train_set._batch_size}{hpstr})"
+            f"batch_size={self.train_set._batch_size}{hpstr})"
 
     def __repr__(self) -> str:
         return str(self)
@@ -380,7 +380,10 @@ def evaluate(self, evaluator: Evaluator, test_set: FastDataLoader) -> dict[str,
             the results.
         """
         if test_set is not None and self.personalized_model is not None:
-            return evaluator.evaluate(self._last_round, self.personalized_model, test_set)
+            return evaluator.evaluate(self._last_round,
+                                      self.personalized_model,
+                                      test_set,
+                                      device=self.device)
         return {}
 
     def state_dict(self) -> dict:

fluke/data/support.py

@@ -50,7 +50,8 @@ def __init__(self, root, train=True, transform=None, target_transform=None, down
 
         data_file = (self.training_file if self.train
                      else self.test_file)
-        self.data, self.targets = torch.load(os.path.join(self.processed_folder, data_file))
+        self.data, self.targets = torch.load(os.path.join(self.processed_folder, data_file),
+                                             weights_only=False)
 
     def __getitem__(self, index):
         """Get images and target for data loader.

fluke/run.py

@@ -91,7 +91,7 @@ def centralized(alg_cfg: str = typer.Argument(..., help='Config file for the alg
             optimizer.step()
         scheduler.step()
 
-        epoch_eval = evaluator.evaluate(e+1, model, test_loader, criterion)
+        epoch_eval = evaluator.evaluate(e+1, model, test_loader, criterion, device=device)
         history.append(epoch_eval)
         for k, v in epoch_eval.items():
             log.add_scalar(k, v, e+1)
@@ -198,7 +198,7 @@ def clients_only(alg_cfg: str = typer.Argument(..., help='Config file for the al
                 optimizer.step()
             scheduler.step()
 
-            client_eval = evaluator.evaluate(e+1, model, test_loader, criterion)
+            client_eval = evaluator.evaluate(e+1, model, test_loader, criterion, device=device)
             running_evals[i].append(client_eval)
 
         log.pretty_log(client_eval, title=f"Client [{i}] Performance")

fluke/server.py

@@ -198,7 +198,7 @@ def evaluate(self, evaluator: Evaluator, test_set: FastDataLoader) -> dict[str,
                 the results.
         """
         if test_set is not None:
-            return evaluator.evaluate(self.rounds + 1, self.model, test_set)
+            return evaluator.evaluate(self.rounds + 1, self.model, test_set, device=self.device)
         return {}
 
     def finalize(self) -> None:

tests/configs/alg/kafe.yaml

@@ -0,0 +1,17 @@
+hyperparameters:
+  client:
+    batch_size: 64
+    local_epochs: 5
+    loss: CrossEntropyLoss
+    optimizer:
+      lr: 0.1
+      momentum: 0.5
+      # weight_decay: 0.0001
+    scheduler:
+      gamma: 1
+      step_size: 1
+  server:
+    weighted: true
+    bandwidth: 1.0
+  model: MNIST_2NN
+name: fluke.algorithms.kafe.Kafe

tests/configs/exp.yaml

@@ -11,6 +11,10 @@ exp:
   seed: 42
 logger:
   name: local
+eval:
+  locals: true
+  pre_fit: true
+  post_fit: true
 protocol:
   eligible_perc: 1
   n_clients: 100

tests/test_alg.py

@@ -399,6 +399,10 @@ def test_fedsgd():
     #                          "./tests/configs/alg/fedsgd.yaml", oncpu=False)
 
 
+def test_kafe():
+    kafe, log = _test_algo("./tests/configs/exp.yaml", "./tests/configs/alg/kafe.yaml")
+
+
 def test_lgfedavg():
     lgfedavg, log = _test_algo("./tests/configs/exp.yaml", "./tests/configs/alg/lg_fedavg.yaml")
     # lgfedavg, log = _test_algo("./tests/configs/exp.yaml",
@@ -451,7 +455,7 @@ def test_superfed():
     # test_fedrep()
     # test_lgfedavg()
     # test_moon()
-    test_fedbn()
+    # test_fedbn()
     # test_pfedme()  # TO BE CHECKED
     # test_scaffold()
     # test_superfed()
@@ -465,3 +469,4 @@ def test_superfed():
     # test_fedavgm()
     # test_fedhp()
     # test_fednh()
+    test_kafe()