all the clients are created and then only a selected fraction is used for the trainings

Author: gaborvecsei

fed_learn/__init__.py

@@ -3,5 +3,5 @@
 from .fed_client import Client
 from .fed_server import Server
 from .models import create_model, set_model_weights
-from .utils import get_rid_of_the_models
+from .utils import get_rid_of_the_models, print_selected_clients
 from .weight_summarizer import FedAvg, WeightSummarizer

fed_learn/args_helper.py

@@ -3,9 +3,11 @@
 
 def get_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("-e", "--global-epochs", help="Number of global (server) epochs", type=int, default=5,
+    parser.add_argument("-e", "--global-epochs", help="Number of global (server) epochs", type=int, default=10,
+                        required=False)
+    parser.add_argument("-c", "--clients", help="Number of clients", type=int, default=100, required=False)
+    parser.add_argument("-f", "--fraction", help="Client fraction to use", type=float, default=0.2,
                         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)

fed_learn/fed_server.py

@@ -9,10 +9,12 @@
 
 class Server:
     def __init__(self, model_fn: Callable,
-                 nb_clients: int,
                  weight_summarizer: WeightSummarizer,
+                 nb_clients: int = 100,
+                 client_fraction: float = 0.2,
                  only_debugging: bool = True):
         self.nb_clients = nb_clients
+        self.client_fraction = client_fraction
         self.weight_summarizer = weight_summarizer
 
         # Initialize the global model's weights
@@ -64,33 +66,34 @@ def _generate_data_indices(self):
     def _get_data_indices_for_client(self, client: int):
         return self.client_data_indices[client]
 
-    def send_train_data(self, client):
+    def _send_train_data_to_client(self, client):
         relevant_data_point_indices = self._get_data_indices_for_client(client.id)
         x = self.x_train[relevant_data_point_indices]
         y = self.y_train[relevant_data_point_indices]
         client.receive_data(x, y)
         return x, y
 
+    def send_train_data(self):
+        self._generate_data_indices()
+        for c in self.clients:
+            self._send_train_data_to_client(c)
+
     def send_model(self, client):
         client.receive_and_init_model(self.model_fn, self.global_model_weights)
 
     def init_for_new_epoch(self):
-        # Reset clients
-        self.clients.clear()
         # Reset the collected weights
         self.client_model_weights.clear()
         # Reset epoch losses
         self.epoch_losses.clear()
-        # Generate new data indices for the clients
-        self._generate_data_indices()
 
     def receive_results(self, client):
         client_weights = client.model.get_weights()
         self.client_model_weights.append(client_weights)
         client.reset_model()
 
     def create_clients(self):
-        # Create new ones
+        # Create all the clients
         for i in range(self.nb_clients):
             client = fed_learn.Client(i)
             self.clients.append(client)
@@ -117,3 +120,10 @@ def test_global_model(self):
         fed_learn.get_rid_of_the_models(model)
 
         return results_dict
+
+    def select_clients(self):
+        nb_clients_to_use = max(int(self.nb_clients * self.client_fraction), 1)
+        client_indices = np.arange(self.nb_clients)
+        np.random.shuffle(client_indices)
+        selected_client_indices = client_indices[:nb_clients_to_use]
+        return np.asarray(self.clients)[selected_client_indices]

fed_learn/utils.py

@@ -1,10 +1,18 @@
-from keras import backend as K
 import gc
+from typing import List
+
+from keras import backend as K
+
+import fed_learn
 
 
 def get_rid_of_the_models(model=None):
-    # TODO: somehow this does not free up the GPU memory (after a while you will get OOM) (tested on Windows 10...)
     K.clear_session()
     if model is not None:
         del model
     gc.collect()
+
+
+def print_selected_clients(clients: List[fed_learn.fed_client.Client]):
+    client_ids = [c.id for c in clients]
+    print("Selected clients for epoch: {0}".format("| ".join(map(str, client_ids))))

fl.py

@@ -7,6 +7,7 @@
 args = fed_learn.get_args()
 
 nb_clients = args.clients
+client_fraction = args.fraction
 nb_global_epochs = args.global_epochs
 debug = args.debug
 
@@ -18,20 +19,22 @@ def model_fn():
 
 
 weight_summarizer = fed_learn.FedAvg()
-server = fed_learn.Server(model_fn, nb_clients, weight_summarizer, debug)
+server = fed_learn.Server(model_fn, weight_summarizer, nb_clients, client_fraction, debug)
 server.update_client_train_params(client_train_params)
+server.create_clients()
+server.send_train_data()
 
 for epoch in range(nb_global_epochs):
     print("Global Epoch {0} is starting".format(epoch))
     server.init_for_new_epoch()
-    server.create_clients()
+    selected_clients = server.select_clients()
 
-    for client in server.clients:
+    fed_learn.print_selected_clients(selected_clients)
+
+    for client in selected_clients:
         print("Client {0} is starting the training".format(client.id))
 
         server.send_model(client)
-        server.send_train_data(client)
-
         hist = client.edge_train(server.get_client_train_param_dict())
         server.epoch_losses.append(hist.history["loss"][-1])
 
@@ -44,9 +47,8 @@ def model_fn():
     print("Loss (client mean): {0}".format(server.global_train_losses[-1]))
 
     global_test_results = server.test_global_model()
-    print("Global test|")
+    print("--- Global test ---")
     for metric_name, value in global_test_results.items():
-        print("_" * 10)
         print("{0}: {1}".format(metric_name, value))
 
     with open("fed_learn_global_test_results.json", 'w') as f: