🐞 fix(method): Fix methods

APFL, Ditto, MetaFed, pFedMe

src/client/apfl.py

@@ -25,6 +25,7 @@ def _re_init(src):
  return deepcopy(target.state_dict())
 
  self.optimizer.add_param_group({"params": trainable_params(self.local_model)})
+ self.init_optimizer_state = deepcopy(self.optimizer.state_dict())
 
  def set_parameters(self, package: dict[str, Any]):
  super().set_parameters(package)

src/client/ditto.py

@@ -11,6 +11,7 @@ def __init__(self, **commons):
  super().__init__(**commons)
  self.pers_model = deepcopy(self.model).to(self.device)
  self.optimizer.add_param_group({"params": trainable_params(self.pers_model)})
+ self.init_optimizer_state = deepcopy(self.optimizer.state_dict())
 
  def set_parameters(self, package: dict[str, Any]):
  super().set_parameters(package)

src/client/fedavg.py

@@ -43,7 +43,9 @@ def __init__(
  self.init_optimizer_state = deepcopy(self.optimizer.state_dict())
  self.lr_scheduler: torch.optim.lr_scheduler.LRScheduler = None
  self.init_lr_scheduler_state: dict = None
+ self.lr_scheduler_cls = None
  if lr_scheduler_cls is not None:
+ self.lr_scheduler_cls = lr_scheduler_cls
  self.lr_scheduler = lr_scheduler_cls(optimizer=self.optimizer)
  self.init_lr_scheduler_state = deepcopy(self.lr_scheduler.state_dict())
 

src/client/metafed.py

@@ -7,7 +7,7 @@
 
 from src.client.fedavg import FedAvgClient
 from torch.utils.data import Subset, DataLoader
-from src.utils.tools import trainable_params, evalutate_model
+from src.utils.tools import evalutate_model
 
 
 class MetaFedClient(FedAvgClient):
@@ -60,12 +60,13 @@ def set_parameters(self, package: dict[str, Any]):
  if package["optimizer_state"]:
  self.optimizer.load_state_dict(package["optimizer_state"])
  else:
- self.optimizer = self.optimizer_cls(params=trainable_params(self.model))
+ self.optimizer.load_state_dict(self.init_optimizer_state)
 
- if package["lr_scheduler_state"]:
- self.lr_scheduler.load_state_dict(package["lr_scheduler_state"])
- elif self.lr_scheduler_cls is not None:
- self.lr_scheduler = self.lr_scheduler_cls(optimizer=self.optimizer)
+ if self.lr_scheduler is not None:
+ if package["lr_scheduler_state"]:
+ self.lr_scheduler.load_state_dict(package["lr_scheduler_state"])
+ else:
+ self.lr_scheduler.load_state_dict(self.init_lr_scheduler_state)
 
  self.client_flag = package["client_flag"]
  self.model.load_state_dict(package["student_model_params"])

src/client/pfedme.py

@@ -20,7 +20,8 @@ def __init__(self, **commons):
  self.args.pfedme.lamda,
  self.args.pfedme.mu,
  )
- self.lr_scheduler = self.lr_scheduler_cls(self.optimizer)
+ if self.lr_scheduler_cls is not None:
+ self.lr_scheduler = self.lr_scheduler_cls(self.optimizer)
 
  def set_parameters(self, package: dict[str, Any]):
  super().set_parameters(package)