Added config option to enable/disable dropout

Signed-off-by: Jim Martens <github@2martens.de>

src/twomartens/masterthesis/cli.py

@@ -185,7 +185,7 @@ def _ssd_train(args: argparse.Namespace) -> None:
     
     conf_obj = conf.Config()
     
-    use_dropout = _ssd_is_dropout(args)
+    use_dropout = _ssd_is_bayesian(args)
     paths = _ssd_train_prepare_paths(args, conf_obj)
     
     ground_truth = _ssd_train_load_gt(conf_obj)
@@ -240,19 +240,16 @@ def _ssd_test(args: argparse.Namespace) -> None:
     _init_eager_mode()
     
     conf_obj = conf.Config()
-    use_dropout = _ssd_is_dropout(args)
+    use_bayesian = _ssd_is_bayesian(args)
     paths = _ssd_test_prepare_paths(args, conf_obj)
     ground_truth = _ssd_test_load_gt(conf_obj)
 
-    ssd_model, predictor_sizes = ssd.get_model(use_dropout,
+    ssd_model, predictor_sizes = ssd.get_model(use_bayesian=use_bayesian,
-                                               keras_ssd300_dropout.ssd_300_dropout,
+                                               bayesian_model=keras_ssd300_dropout.ssd_300_dropout,
-                                               keras_ssd300.ssd_300,
+                                               vanilla_model=keras_ssd300.ssd_300,
-                                               conf_obj.parameters.ssd_image_size,
+                                               conf_obj=conf_obj,
-                                               conf_obj.parameters.nr_classes,
+                                               mode="training",
-                                               "training",
+                                               pre_trained_weights_file=paths.weights_file)
-                                               conf_obj.parameters.ssd_dropout_rate,
-                                               conf_obj.parameters.ssd_top_k,
-                                               paths.weights_file)
 
     loss_func = ssd.get_loss_func()
     ssd.compile_model(ssd_model, conf_obj.parameters.learning_rate, loss_func)
@@ -270,7 +267,7 @@ def _ssd_test(args: argparse.Namespace) -> None:
                 model=ssd_model,
                 conf_obj=conf_obj,
                 steps_per_epoch=steps_per_epoch,
-                use_dropout=use_dropout,
+                use_bayesian=use_bayesian,
                 nr_digits=nr_digits,
                 output_path=paths.output_path)
 
@@ -730,7 +727,7 @@ def _visualise_get_config_values(config_get: Callable[[str], Union[str, int, flo
     return output_path, coco_path, ground_truth_path
 
 
-def _ssd_is_dropout(args: argparse.Namespace) -> bool:
+def _ssd_is_bayesian(args: argparse.Namespace) -> bool:
     return False if args.network == "ssd" else True
 
 

src/twomartens/masterthesis/config.py

@@ -68,6 +68,7 @@ _CONFIG_PROPS = {
         "ssd_entropy_threshold_max": (float, "2.5"),
         "ssd_test_pretrained": (bool, "False"),
         "ssd_use_coco": (bool, "False"),
+        "ssd_use_dropout": (bool, "True"),
         "nr_trajectories": (int, "-1")
     }
 }

src/twomartens/masterthesis/ssd.py

@@ -50,53 +50,48 @@ K = tf.keras.backend
 tfe = tf.contrib.eager
 
 
-def get_model(use_dropout: bool,
+def get_model(use_bayesian: bool,
-              dropout_model: callable, vanilla_model: callable,
+              bayesian_model: callable, vanilla_model: callable,
-              image_size: int, nr_classes: int, mode: str,
+              conf_obj: config.Config,
-              dropout_rate: float, top_k: int,
+              mode: str,
-              pre_trained_weights_file: Optional[str] = None,
+              pre_trained_weights_file: Optional[str] = None) -> Tuple[tf.keras.models.Model, np.ndarray]:
-              iou_threshold: Optional[float] = None) -> Tuple[tf.keras.models.Model, np.ndarray]:
     """
     Returns the correct SSD model and the corresponding predictor sizes.
     
     Args:
-        use_dropout: True if dropout variant should be used, False otherwise
+        use_bayesian: True if Bayesian variant should be used, False otherwise
-        dropout_model: function to build dropout SSD model
+        bayesian_model: function to build Bayesian SSD model
         vanilla_model: function to build vanilla SSD model
-        image_size: size of the resized images
+        conf_obj: configuration object
-        nr_classes: number of classes
         mode: one of "training", "inference", "inference_fast"
-        dropout_rate: rate for dropout layers (only applies if dropout is used)
-        top_k: number of highest scoring predictions kept for each batch item
         pre_trained_weights_file: path to h5 file with pre-trained weights
-        iou_threshold: all boxes with higher iou to local maximum box are suppressed
-                       (only relevant for inference modes)
 
     Returns:
         SSD model, predictor_sizes
     """
     
-    image_size = (image_size, image_size, 3)
+    image_size = (conf_obj.parameters.ssd_image_size, conf_obj.parameters.ssd_image_size, 3)
     scales = [0.07, 0.15, 0.33, 0.51, 0.69, 0.87, 1.05]
-    if use_dropout:
+    if use_bayesian:
-        model, predictor_sizes = dropout_model(
+        model, predictor_sizes = bayesian_model(
             image_size=image_size,
-            n_classes=nr_classes,
+            n_classes=conf_obj.parameters.nr_classes,
             mode=mode,
-            iou_threshold=iou_threshold,
+            iou_threshold=conf_obj.parameters.ssd_iou_threshold,
-            dropout_rate=dropout_rate,
+            dropout_rate=conf_obj.parameters.ssd_dropout_rate,
-            top_k=top_k,
+            top_k=conf_obj.parameters.ssd_top_k,
             scales=scales,
             return_predictor_sizes=True,
-            coords="corners"
+            coords="corners",
+            use_dropout=conf_obj.parameters.ssd_use_dropout
         )
     else:
         model, predictor_sizes = vanilla_model(
             image_size=image_size,
-            n_classes=nr_classes,
+            n_classes=conf_obj.parameters.nr_classes,
             mode=mode,
-            iou_threshold=iou_threshold,
+            iou_threshold=conf_obj.parameters.ssd_iou_threshold,
-            top_k=top_k,
+            top_k=conf_obj.parameters.ssd_top_k,
             scales=scales,
             return_predictor_sizes=True,
             coords="corners"
@@ -150,7 +145,7 @@ def predict(generator: callable,
             model: tf.keras.models.Model,
             conf_obj: config.Config,
             steps_per_epoch: int,
-            use_dropout: bool,
+            use_bayesian: bool,
             nr_digits: int,
             output_path: str) -> None:
     """
@@ -163,13 +158,13 @@ def predict(generator: callable,
         model: compiled and trained Keras model
         conf_obj: configuration object
         steps_per_epoch: number of batches per epoch
-        use_dropout: if True, multiple forward passes and observations will be used
+        use_bayesian: if True, multiple forward passes and observations will be used
         nr_digits: number of digits needed to print largest batch number
         output_path: the path in which the results should be saved
     """
-    output_file, label_output_file = _predict_prepare_paths(output_path, use_dropout)
+    output_file, label_output_file = _predict_prepare_paths(output_path, use_bayesian)
     
-    _predict_loop(generator, use_dropout, steps_per_epoch,
+    _predict_loop(generator, use_bayesian, steps_per_epoch,
                   dropout_step=functools.partial(
                       _predict_dropout_step,
                       model=model,