[in progress] argparse approach

libs/scoring.py

@@ -26,7 +26,7 @@ def plot_confusion_matrix(y_true, y_pred, classes,
                           normalize=True,
                           title=None,
                           cmap=plt.cm.Blues,
-                          savedir="score_files"):
+                          savedir="predictions"):
     """
     This function prints and plots the confusion matrix.
     Normalization can be applied by setting `normalize=True`.
@@ -79,8 +79,7 @@ def plot_confusion_matrix(y_true, y_pred, classes,
     # save to directory
     if not os.path.isdir(savedir):
         os.mkdir(savedir)
-    savefile = savedir + '/score-' + title
-
+    savefile = title
     plt.savefig(savefile)
     return savefile, cm
 
@@ -119,7 +118,7 @@ def score_masks(labelfile, predictionfile):
 
     return precision, recall, f1, savefile
 
-def score_predictions(dataset, basedir='.'):
+def score_predictions(dataset, basedir='predictions'):
 
     scores = []
 

libs/training.py

@@ -14,11 +14,11 @@
 from wandb.fastai import WandbCallback
 
 
-def train_model(dataset):
+def train_model(dataset, config):
     """ Trains a DynamicUnet on the dataset """
 
-    epochs = 15
-    lr     = 1e-4
+    epochs = config["epochs"]
+    lr     = config["lr"]
     size   = 300
     wd     = 1e-2
     bs     = 8 # reduce this if you are running out of GPU memory

main.py

@@ -3,33 +3,89 @@
 from libs import training
 from libs import datasets
 
+import argparse
 import wandb
 
-if __name__ == '__main__':
-
-    dataset = 'dataset-sample'  #  0.5 GB download
-    # dataset = 'dataset-medium' # 9.0 GB download
+# config settings/hyperparameters
+# these defaults can be edited here or overwritten via command line
+# note, other config defaults for training and inference are defined within
+# relevant files under "libs/"
+MODEL_NAME = ""
+DATASET_NAME = "sample"
+EPOCHS = 15
+LEARNING_RATE = 1e-4
 
+def train_model(args, dataset_name):
     config = {
-        'name' : 'baseline',
-        'dataset' : dataset,
+        'dataset' : dataset_name,
+        'epochs' : args.epochs,
+        'lr' : args.learning_rate
     }
 
-    wandb.init(config=config)
+    wandb.init(config=config, name=args.model_name)
 
-    datasets.download_dataset(dataset)
+    datasets.download_dataset(args.dataset)
 
     # train the baseline model and save it in models folder
-    training.train_model(dataset)
+    training.train_model(dataset, config)
 
     # use the train model to run inference on all test scenes
     inference.run_inference(dataset)
 
     # scores all the test images compared to the ground truth labels then
     # send the scores (f1, precision, recall) and prediction images to wandb
     score, predictions = scoring.score_predictions(dataset)
-    wandb.config.update(score)
+    print(score)
+    wandb.log(score)
 
     for f1, f2 in predictions:
-        wandb.save( f1 )
-        wandb.save( f2 )
+        wandb.save(f1)
+        wandb.save(f2)   
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        "-m",
+        "--model_name",
+        type=str,
+        default=MODEL_NAME,
+        help="Name of this model/run")
+    parser.add_argument(
+        "-d",
+        "--dataset_name",
+        type=str,
+        default=DATASET_NAME,
+        help="Dataset name (current options: sample or medium)")
+    parser.add_argument(
+        "-e",
+        "--epochs",
+        type=int,
+        default=EPOCHS,
+        help="Number of training epochs")
+    parser.add_argument(
+        "-lr",
+        "--learning_rate",
+        type=float,
+        default=LEARNING_RATE,
+        help="Learning rate")
+    parser.add_argument(
+        "-q",
+        "--dry_run",
+        action="store_true",
+        help="Dry run (do not log to wandb)")
+    args = parser.parse_args()
+
+    # easier testing--don't log to wandb if dry run is set
+    if args.dry_run:
+        os.environ['WANDB_MODE'] = 'dryrun'
+
+    dataset = ""
+    if args.dataset_name == "sample": 
+        dataset = 'dataset-sample'  #  0.5 GB download
+    elif args.dataset_name == "medium":
+        dataset = 'dataset-medium' # 9.0 GB download
+    if not dataset:
+        print("Unable to load dataset (valid options are 'sample' and 'medium'): ", args.dataset_name)
+        exit(0)
+
+    train_model(args, dataset_name)

main_keras.py

@@ -8,7 +8,7 @@
 
 if __name__ == '__main__':
     dataset = 'dataset-sample'  #  0.5 GB download
-    # dataset = 'dataset-medium' # 9.0 GB download
+    #dataset = 'dataset-medium' # 9.0 GB download
 
     config = {
         'name' : 'baseline-keras',
@@ -29,4 +29,5 @@
     # scores all the test images compared to the ground truth labels then
     # send the scores (f1, precision, recall) and prediction images to wandb
     score, _ = scoring.score_predictions(dataset, basedir=wandb.run.dir)
-    wandb.config.update(score)
+    print(score)
+    wandb.log(score)

requirements.txt

@@ -5,6 +5,7 @@ numpy==1.17.1
 opencv_python==3.4.3.18
 sklearn
 tensorflow
+tensorflow-gpu
 torch
 typing==3.6.6
 wandb