epoch, clients, debug arguments added

Author: gaborvecsei

fl.py

@@ -1,25 +1,36 @@
+import argparse
+
 import fed_learn
+import numpy as np
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-ge", "--global-epochs", help="Number of global (server) epochs", type=int, default=5,
+                    required=False)
+parser.add_argument("-c", "--clients", help="Number of clients", type=int, default=10, required=False)
+parser.add_argument("-d", "--debug", help="Debugging", action="store_true", required=False)
+args = parser.parse_args()
 
-NB_CLIENTS = 10
-NB_EPOCHS = 3
+nb_clients = args.clients
+nb_epochs = args.global_epochs
+debug = args.debug
 
 
 def model_fn():
     return fed_learn.create_model((32, 32, 3), 10)
 
 
 weight_summarizer = fed_learn.FedAvg()
-server = fed_learn.Server(model_fn, NB_CLIENTS, weight_summarizer)
+server = fed_learn.Server(model_fn, nb_clients, weight_summarizer, debug)
 
-for epoch in range(NB_EPOCHS):
+for epoch in range(nb_epochs):
     print("Global Epoch {0} is starting".format(epoch))
     server.create_clients()
     server.init_for_new_epoch()
 
     loss = []
 
     for client in server.clients:
-        print("Client {0} is starting the training".format({client.id}))
+        print("Client {0} is starting the training".format(client.id))
 
         server.send_model(client)
         server.send_train_data(client)
@@ -30,6 +41,7 @@ def model_fn():
         server.receive_results(client)
 
     server.summarize_weights()
+    print("Loss (mean): {0}".format(np.mean(loss)))
     print("-" * 30)
 
     # TODO: test the base model with the aggregated weights