client train params are updateable by args and ars are moved to a helper script

Author: gaborvecsei

fed_learn/__init__.py

@@ -1,6 +1,7 @@
+from .args_helper import get_args
 from .data_sampling import iid_data_indices, non_iid_data_indices
 from .fed_client import Client
 from .fed_server import Server
 from .models import create_model, set_model_weights
-from .weight_summarizer import FedAvg, WeightSummarizer
 from .utils import get_rid_of_the_models
+from .weight_summarizer import FedAvg, WeightSummarizer

fed_learn/args_helper.py

@@ -0,0 +1,16 @@
+import argparse
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-e", "--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)
+
+    parser.add_argument("-lr", "--learning-rate", help="Learning rate", type=float, default=0.01, required=False)
+    parser.add_argument("-b", "--batch-size", help="Batch Size", type=int, default=32, required=False)
+    parser.add_argument("-ce", "--client-epochs", help="Number of epochs for the clients", type=int, default=5,
+                        required=False)
+    args = parser.parse_args()
+    return args

fed_learn/fed_server.py

@@ -53,15 +53,10 @@ def __init__(self, model_fn: Callable,
         self.client_model_weights = []
 
         # Training parameters used by the clients
-        # TODO: this should be configurable
-        self.train_dict = {"batch_size": 32,
-                           "epochs": 5,
-                           "verbose": 1,
-                           "shuffle": True}
-
-        if only_debugging:
-            # TODO: remove me
-            self.train_dict["epochs"] = 1
+        self.client_train_params_dict = {"batch_size": 32,
+                                         "epochs": 5,
+                                         "verbose": 1,
+                                         "shuffle": True}
 
     def _generate_data_indices(self):
         self.client_data_indices = fed_learn.iid_data_indices(self.nb_clients, len(self.x_train))
@@ -105,7 +100,10 @@ def summarize_weights(self):
         self.global_model_weights = new_weights
 
     def get_client_train_param_dict(self):
-        return self.train_dict
+        return self.client_train_params_dict
+
+    def update_client_train_params(self, param_dict: dict):
+        self.client_train_params_dict.update(param_dict)
 
     def test_global_model(self):
         model = self.model_fn()

fed_learn/models.py

@@ -3,7 +3,7 @@
 from keras.applications.vgg16 import VGG16
 
 
-def create_model(input_shape: tuple, nb_classes: int, init_with_imagenet: bool = False):
+def create_model(input_shape: tuple, nb_classes: int, init_with_imagenet: bool = False, learning_rate: float = 0.01):
     weights = None
     if init_with_imagenet:
         weights = "imagenet"
@@ -21,7 +21,7 @@ def create_model(input_shape: tuple, nb_classes: int, init_with_imagenet: bool =
     model = models.Model(model.input, x)
 
     loss = losses.categorical_crossentropy
-    optimizer = optimizers.Adam(lr=0.001)
+    optimizer = optimizers.SGD(lr=learning_rate, decay=0.99)
 
     model.compile(optimizer, loss, metrics=["accuracy"])
     return model

fl.py

@@ -1,30 +1,27 @@
-import argparse
 import json
 
 import numpy as np
 
 import fed_learn
 
-parser = argparse.ArgumentParser()
-parser.add_argument("-e", "--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()
+args = fed_learn.get_args()
 
 nb_clients = args.clients
-nb_epochs = args.global_epochs
+nb_global_epochs = args.global_epochs
 debug = args.debug
 
+client_train_params = {"epochs": args.client_epochs, "batch_size": args.batch_size}
+
 
 def model_fn():
-    return fed_learn.create_model((32, 32, 3), 10, init_with_imagenet=False)
+    return fed_learn.create_model((32, 32, 3), 10, init_with_imagenet=False, learning_rate=args.learning_rate)
 
 
 weight_summarizer = fed_learn.FedAvg()
 server = fed_learn.Server(model_fn, nb_clients, weight_summarizer, debug)
+server.update_client_train_params(client_train_params)
 
-for epoch in range(nb_epochs):
+for epoch in range(nb_global_epochs):
     print("Global Epoch {0} is starting".format(epoch))
     server.init_for_new_epoch()
     server.create_clients()