files are stored in the experiments folder (#4)

Author: gaborvecsei

.gitignore

@@ -1,3 +1,4 @@
 .idea/
 __pycache__
+experiments/
 *.json

fed_learn/__init__.py

@@ -1,4 +1,4 @@
-from .args_helper import get_args
+from .args_helper import get_args, save_args_as_json
 from .data_sampling import iid_data_indices, non_iid_data_indices
 from .fed_client import Client
 from .fed_server import Server

fed_learn/args_helper.py

@@ -1,8 +1,12 @@
 import argparse
+import json
 
 
 def get_args():
     parser = argparse.ArgumentParser()
+    parser.add_argument("-n", "--name", help="Name of the experiment", type=str, required=True)
+    parser.add_argument("-oe", "--overwrite-experiment", help="Overwrite existing experiment", action="store_true",
+                        required=False)
     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)
@@ -16,3 +20,8 @@ def get_args():
                         required=False)
     args = parser.parse_args()
     return args
+
+
+def save_args_as_json(args, path):
+    with open(str(path), "w") as f:
+        json.dump(args.__dict__, f, sort_keys=True, indent=4)

federated_learning.py

@@ -1,15 +1,19 @@
 import json
+from pathlib import Path
 
 import numpy as np
 
 import fed_learn
 
 args = fed_learn.get_args()
 
-nb_clients = args.clients
-client_fraction = args.fraction
-nb_global_epochs = args.global_epochs
-debug = args.debug
+EXPERIMENT_FOLDER_PATH = Path(__file__).resolve().parent / "experiments" / args.name
+EXPERIMENT_FOLDER_PATH.mkdir(parents=True, exist_ok=args.overwrite_experiment)
+
+args_json_path = EXPERIMENT_FOLDER_PATH / "args.json"
+fed_learn.save_args_as_json(args, EXPERIMENT_FOLDER_PATH / args_json_path)
+
+train_hist_path = EXPERIMENT_FOLDER_PATH / "fed_learn_global_test_results.json"
 
 client_train_params = {"epochs": args.client_epochs, "batch_size": args.batch_size}
 
@@ -19,12 +23,12 @@ def model_fn():
 
 
 weight_summarizer = fed_learn.FedAvg()
-server = fed_learn.Server(model_fn, weight_summarizer, nb_clients, client_fraction, debug)
+server = fed_learn.Server(model_fn, weight_summarizer, args.clients, args.fraction, args.debug)
 server.update_client_train_params(client_train_params)
 server.create_clients()
 server.send_train_data()
 
-for epoch in range(nb_global_epochs):
+for epoch in range(args.global_epochs):
     print("Global Epoch {0} is starting".format(epoch))
     server.init_for_new_epoch()
     selected_clients = server.select_clients()
@@ -51,7 +55,7 @@ def model_fn():
     for metric_name, value in global_test_results.items():
         print("{0}: {1}".format(metric_name, value))
 
-    with open("fed_learn_global_test_results.json", 'w') as f:
+    with open(str(train_hist_path), 'w') as f:
         json.dump(server.global_test_metrics_dict, f)
 
     print("_" * 30)