Add calculate manual grad and update manual grad of parameter each layer

client.py

@@ -7,14 +7,14 @@
 
 import src.Log
 from src.RpcClient import RpcClient
-from src.Scheduler import Scheduler
+from src.Schedule_zb import Scheduler
 
 
 parser = argparse.ArgumentParser(description="Split learning framework")
 parser.add_argument('--layer_id', type=int, required=True, help='ID of layer, start from 1')
 parser.add_argument('--device', type=str, required=False, help='Device of client')
 parser.add_argument('--event_time', type=bool, default=False, required=False, help='Log event time for debug mode')
-parser.add_argument('--performance', type=int, required=False, help='Cluster by device')
+parser.add_argument('--p', type=int, required=False, help='Cluster by device')
 
 args = parser.parse_args()
 
@@ -45,15 +45,15 @@
 connection = pika.BlockingConnection(pika.ConnectionParameters(address, 5672, f'{virtual_host}', credentials))
 channel = connection.channel()
 
-if args.performance is None:
+if args.p is None:
     performance = -1
 else:
-    performance = args.performance
+    performance = args.p
 
 if __name__ == "__main__":
     src.Log.print_with_color("[>>>] Client sending registration message to server...", "red")
     data = {"action": "REGISTER", "client_id": client_id, "layer_id": args.layer_id, "performance": performance, "message": "Hello from Client!"}
-    scheduler = Scheduler(client_id, args.layer_id, channel, device, args.event_time)
+    scheduler = Scheduler(client_id, args.layer_id, channel, device)
     client = RpcClient(client_id, args.layer_id, address, username, password, scheduler.train_on_device, device)
     client.send_to_server(data)
     client.wait_response()

config.yaml

@@ -35,6 +35,9 @@ server:
     AffinityPropagation:
       damping: 0.9
       max_iter: 1000
+  fbw:
+    enable: True
+    chunk: 1
 
 rabbit:
   address: 127.0.0.1

src/RpcClient.py

@@ -74,12 +74,14 @@ def response_message(self, body):
             label_count = self.response['label_count']
             num_layers = self.response['num_layers']
             clip_grad_norm = self.response['clip_grad_norm']
+            chunk = self.response['chunk']
             if self.label_count is None:
                 self.label_count = label_count
             if self.response['cluster'] is not None:
                 self.cluster = self.response['cluster']
             if self.label_count is not None:
                 src.Log.print_with_color(f"Label distribution of client: {self.label_count}", "yellow")
+
             if self.model is None:
                 klass = getattr(src.Model, model_name)
                 full_model = klass()
@@ -116,11 +118,11 @@ def response_message(self, body):
                 subset = torch.utils.data.Subset(self.train_set, selected_indices)
                 train_loader = torch.utils.data.DataLoader(subset, batch_size=batch_size, shuffle=True)
                 if cut_layers[1] != 0:
-                    result, size = self.train_func(self.model, self.global_model, self.label_count, lr, momentum, clip_grad_norm, compute_loss, num_layers, control_count, train_loader, self.cluster, special, alone_train=False)
+                    result, size = self.train_func(self.model, self.global_model, self.label_count, lr, momentum, clip_grad_norm, compute_loss, num_layers, control_count, train_loader, self.cluster, special, alone_train=False, chunk=chunk)
                 else:
-                    result, size = self.train_func(self.model, self.global_model, self.label_count, lr, momentum, clip_grad_norm, compute_loss, num_layers, control_count, train_loader, self.cluster, special, alone_train=True)
+                    result, size = self.train_func(self.model, self.global_model, self.label_count, lr, momentum, clip_grad_norm, compute_loss, num_layers, control_count, train_loader, self.cluster, special, alone_train=True, chunk=chunk)
             else:
-                result, size = self.train_func(self.model, self.global_model, self.label_count, lr, momentum, clip_grad_norm, compute_loss, num_layers, control_count, None, self.cluster, special)
+                result, size = self.train_func(self.model, self.global_model, self.label_count, lr, momentum, clip_grad_norm, compute_loss, num_layers, control_count, None, self.cluster, special, chunk=chunk)
 
             # Stop training, then send parameters to server
             model_state_dict = self.model.state_dict()
@@ -132,6 +134,7 @@ def response_message(self, body):
                     "message": "Sent parameters to Server", "parameters": model_state_dict}
             src.Log.print_with_color("[>>>] Client sent parameters to server", "red")
             self.send_to_server(data)
+            self.model = None
             return True
         elif action == "STOP":
             return False