Fix Scaffold

fedlab/contrib/algorithm/scaffold.py

@@ -74,19 +74,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
+
                 self.optimizer.step()
 
         dy = self.model_parameters - frz_model
         dc = -1.0 / (self.epochs * len(train_loader) * self.lr) * dy - global_c
         self.cs[id] += dc
         return [dy, dc]
+

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])
+            super().set_model(self.parameters[id])