working GPU is configurable

Author: gaborvecsei

fed_learn/args_helper.py

@@ -19,6 +19,7 @@ def get_args():
     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=1,
                         required=False)
+    parser.add_argument("-g", "--gpu", help="GPU to use (-1 is CPU)", type=int, default=0, required=False)
     args = parser.parse_args()
     return args
 

federated_learning.py

@@ -1,4 +1,5 @@
 import json
+import os
 from pathlib import Path
 
 import numpy as np
@@ -7,6 +8,9 @@
 
 args = fed_learn.get_args()
 
+os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
+
 EXPERIMENT_FOLDER_PATH = Path(__file__).resolve().parent / "experiments" / args.name
 EXPERIMENT_FOLDER_PATH.mkdir(parents=True, exist_ok=args.overwrite_experiment)