From cd2d681948a4ad3b9563bd3a95d0fe40e3cddbff Mon Sep 17 00:00:00 2001 From: hua-zi <83271073+hua-zi@users.noreply.github.com> Date: Fri, 3 Nov 2023 21:58:57 +0800 Subject: [PATCH 1/4] Update scaffold.py --- fedlab/contrib/algorithm/scaffold.py | 31 +++++++++++++++++++++------- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/fedlab/contrib/algorithm/scaffold.py b/fedlab/contrib/algorithm/scaffold.py index 56040184..554e45ba 100644 --- a/fedlab/contrib/algorithm/scaffold.py +++ b/fedlab/contrib/algorithm/scaffold.py @@ -19,7 +19,8 @@ def downlink_package(self): def setup_optim(self, lr): self.lr = lr - self.global_c = torch.zeros_like(self.model_parameters) + # self.global_c = torch.zeros_like(self.model_parameters) + self.global_c = torch.zeros_like(self.model_grad_parameters) def global_update(self, buffer): # unpack @@ -74,19 +75,33 @@ def train(self, id, model_parameters, global_c, train_loader): self.optimizer.zero_grad() loss.backward() - grad = self.model_gradients + # grad = self.model_gradients + grad = self.model_grads grad = grad - self.cs[id] + global_c idx = 0 - for parameter in self._model.parameters(): - layer_size = parameter.grad.numel() - shape = parameter.grad.shape - #parameter.grad = parameter.grad - self.cs[id][idx:idx + layer_size].view(parameter.grad.shape) + global_c[idx:idx + layer_size].view(parameter.grad.shape) - parameter.grad.data[:] = grad[idx:idx+layer_size].view(shape)[:] + + parameters = self._model.parameters() + for p in self._model.state_dict().values(): + if p.grad is None: # batchnorm have no grad + layer_size = p.numel() + else: + parameter = next(parameters) + layer_size = parameter.data.numel() + shape = parameter.grad.shape + parameter.grad.data[:] = grad[idx:idx+layer_size].view(shape)[:] idx += layer_size + # for parameter in self._model.parameters(): + # layer_size = parameter.grad.numel() + # shape = parameter.grad.shape + # #parameter.grad = parameter.grad - self.cs[id][idx:idx + layer_size].view(parameter.grad.shape) + global_c[idx:idx + layer_size].view(parameter.grad.shape) + # parameter.grad.data[:] = grad[idx:idx+layer_size].view(shape)[:] + # idx += layer_size + # print(idx) + self.optimizer.step() dy = self.model_parameters - frz_model - dc = -1.0 / (self.epochs * len(train_loader) * self.lr) * dy - global_c + dc = -1.0 / (self.epochs * len(train_loader) * self.lr) * dy - global_c # ???? self.cs[id] += dc return [dy, dc] From e9737992f2a5bb1cacfbb3047499ad876d543925 Mon Sep 17 00:00:00 2001 From: hua-zi <83271073+hua-zi@users.noreply.github.com> Date: Fri, 3 Nov 2023 22:18:29 +0800 Subject: [PATCH 2/4] Update model_maintainer.py --- fedlab/core/model_maintainer.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/fedlab/core/model_maintainer.py b/fedlab/core/model_maintainer.py index 89d1e8b0..836e3c1f 100644 --- a/fedlab/core/model_maintainer.py +++ b/fedlab/core/model_maintainer.py @@ -60,6 +60,20 @@ def model_parameters(self) -> torch.Tensor: """Return serialized model parameters.""" return SerializationTool.serialize_model(self._model) + @property + def model_grads(self) -> torch.Tensor: + """Return serialized model gradients(base on model.state_dict(), Shape is the same as model_parameters).""" + params = self._model.state_dict() + for name, p in self._model.named_parameters(): + params[name].grad = p.grad + for key in params: + if params[key].grad is None: + params[key].grad = torch.zeros_like(params[key]) + gradients = [param.grad.data.view(-1) for param in params.values()] + m_gradients = torch.cat(gradients) + m_gradients = m_gradients.cpu() + return m_gradients + @property def model_gradients(self) -> torch.Tensor: """Return serialized model gradients.""" @@ -117,4 +131,4 @@ def set_model(self, parameters: torch.Tensor = None, id: int = None): if id is None: super().set_model(parameters) else: - super().set_model(self.parameters[id]) \ No newline at end of file + super().set_model(self.parameters[id]) From 2b6d2d8f3f422bebb7eb1348d2bbac6b3a49f60b Mon Sep 17 00:00:00 2001 From: hua-zi <83271073+hua-zi@users.noreply.github.com> Date: Fri, 3 Nov 2023 23:47:06 +0800 Subject: [PATCH 3/4] Update scaffold.py --- fedlab/contrib/algorithm/scaffold.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/fedlab/contrib/algorithm/scaffold.py b/fedlab/contrib/algorithm/scaffold.py index 554e45ba..711c9e80 100644 --- a/fedlab/contrib/algorithm/scaffold.py +++ b/fedlab/contrib/algorithm/scaffold.py @@ -19,8 +19,7 @@ def downlink_package(self): def setup_optim(self, lr): self.lr = lr - # self.global_c = torch.zeros_like(self.model_parameters) - self.global_c = torch.zeros_like(self.model_grad_parameters) + self.global_c = torch.zeros_like(self.model_parameters) def global_update(self, buffer): # unpack @@ -79,10 +78,10 @@ def train(self, id, model_parameters, global_c, train_loader): grad = self.model_grads grad = grad - self.cs[id] + global_c idx = 0 - + parameters = self._model.parameters() for p in self._model.state_dict().values(): - if p.grad is None: # batchnorm have no grad + if p.grad is None: # Batchnorm have no grad layer_size = p.numel() else: parameter = next(parameters) @@ -97,11 +96,11 @@ def train(self, id, model_parameters, global_c, train_loader): # #parameter.grad = parameter.grad - self.cs[id][idx:idx + layer_size].view(parameter.grad.shape) + global_c[idx:idx + layer_size].view(parameter.grad.shape) # parameter.grad.data[:] = grad[idx:idx+layer_size].view(shape)[:] # idx += layer_size - # print(idx) self.optimizer.step() dy = self.model_parameters - frz_model - dc = -1.0 / (self.epochs * len(train_loader) * self.lr) * dy - global_c # ???? + dc = -1.0 / (self.epochs * len(train_loader) * self.lr) * dy - global_c self.cs[id] += dc return [dy, dc] + From f8bf0673abe1b9ca5b3ab868f4d67be176d77c14 Mon Sep 17 00:00:00 2001 From: hua-zi <83271073+hua-zi@users.noreply.github.com> Date: Fri, 3 Nov 2023 23:47:42 +0800 Subject: [PATCH 4/4] Update scaffold.py