From 01f43651b5f86944f5346b4889975915cf9cd928 Mon Sep 17 00:00:00 2001 From: Jim Martens Date: Fri, 13 Sep 2019 12:35:52 +0200 Subject: [PATCH] Fixed ssd train function Signed-off-by: Jim Martens --- src/twomartens/masterthesis/cli.py | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/src/twomartens/masterthesis/cli.py b/src/twomartens/masterthesis/cli.py index 538f500..30c7e1a 100644 --- a/src/twomartens/masterthesis/cli.py +++ b/src/twomartens/masterthesis/cli.py @@ -185,20 +185,17 @@ def _ssd_train(args: argparse.Namespace) -> None: conf_obj = conf.Config() - use_dropout = _ssd_is_bayesian(args) + use_bayesian = _ssd_is_bayesian(args) paths = _ssd_train_prepare_paths(args, conf_obj) ground_truth = _ssd_train_load_gt(conf_obj) - ssd_model, predictor_sizes = ssd.get_model(use_dropout, - keras_ssd300_dropout.ssd_300_dropout, - keras_ssd300.ssd_300, - conf_obj.parameters.ssd_image_size, - conf_obj.parameters.nr_classes, - "training", - conf_obj.parameters.ssd_dropout_rate, - conf_obj.parameters.ssd_top_k, - paths.pre_trained_weights_file) + ssd_model, predictor_sizes = ssd.get_model(use_bayesian=use_bayesian, + bayesian_model=keras_ssd300_dropout.ssd_300_dropout, + vanilla_model=keras_ssd300.ssd_300, + conf_obj=conf_obj, + mode="training", + pre_trained_weights_file=paths.pre_trained_weights_file) loss_func = ssd.get_loss_func() ssd.compile_model(ssd_model, conf_obj.parameters.learning_rate, loss_func)