serverper.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. from flcore.clients.clientper import clientPer
  2. from flcore.servers.serverbase import Server
  3. import copy
  4. import os
  5. import logging
  6. class FedPer(Server):
  7. def __init__(self, args, times):
  8. super().__init__(args, times)
  9. self.message_hp = f"{args.algorithm}, lr:{args.local_learning_rate:.5f}"
  10. clientObj = clientPer
  11. self.message_hp_dash = self.message_hp.replace(", ", "-")
  12. self.hist_result_fn = os.path.join(args.hist_dir, f"{self.actual_dataset}-{self.message_hp_dash}-{args.goal}-{self.times}.h5")
  13. self.set_clients(args, clientObj)
  14. print(f"\nJoin ratio / total clients: {self.join_ratio} / {self.num_clients}")
  15. print("Finished creating server and clients.")
  16. # self.load_model()
  17. def train(self):
  18. for i in range(self.global_rounds):
  19. self.selected_clients = self.select_clients()
  20. self.send_models()
  21. print(f"\n------------- Round number: [{i+1:3d}/{self.global_rounds}]-------------")
  22. print(f"==> Training for {len(self.selected_clients)} clients...", flush=True)
  23. for client in self.selected_clients:
  24. client.train()
  25. self.receive_models()
  26. self.aggregate_parameters()
  27. if i%self.eval_gap == 0:
  28. print("==> Evaluating personalized models...", flush=True)
  29. self.send_models(mode="all")
  30. self.evaluate(self.global_model)
  31. print(f"==> Best mean personalized accuracy: {self.best_mean_test_acc*100:.2f}%", flush=True)
  32. self.save_results(fn=self.hist_result_fn)
  33. message_res = f"\ttest_acc:{self.best_mean_test_acc:.6f}"
  34. logging.info(self.message_hp + message_res)
  35. def receive_models(self):
  36. assert (len(self.selected_clients) > 0)
  37. self.uploaded_weights = []
  38. tot_samples = 0
  39. self.uploaded_ids = []
  40. self.uploaded_models = []
  41. for client in self.selected_clients:
  42. self.uploaded_weights.append(client.train_samples)
  43. tot_samples += client.train_samples
  44. self.uploaded_ids.append(client.id)
  45. self.uploaded_models.append(client.model.base)
  46. for i, w in enumerate(self.uploaded_weights):
  47. self.uploaded_weights[i] = w / tot_samples
  48. def prepare_global_model(self):
  49. temp_model = copy.deepcopy(self.global_model) # base
  50. self.global_model = copy.deepcopy(self.clients[0].model)
  51. for p_t, p_g in zip(temp_model.parameters(), self.global_model.base.parameters()):
  52. p_g.data = p_t.data.clone()
  53. for p in self.global_model.predictor.parameters():
  54. p.data.zero_()
  55. for c in self.clients:
  56. for p_g, p_c in zip(self.global_model.predictor.parameters(), c.model.predictor.parameters()):
  57. p_g.data += p_c.data * c.train_samples
  58. return