From a785cac631cc8263a8f3d9d0fc51411c590ac5f2 Mon Sep 17 00:00:00 2001 From: Tim Koornstra Date: Wed, 16 Aug 2023 11:32:54 +0200 Subject: [PATCH] Last dense layer output not hardcoded anymore --- src/main.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/main.py b/src/main.py index b64669af..9ec07dd6 100644 --- a/src/main.py +++ b/src/main.py @@ -525,10 +525,19 @@ def main(): use_mask=args.use_mask ) training_generator, validation_generator, test_generator, inference_generator, utilsObject, train_batches = loader.generators() + + # Get the prediction model by taking the last dense layer of the full + # model + last_dense_layer = None + for layer in reversed(model.layers): + if layer.name.startswith('dense'): + last_dense_layer = layer + break + prediction_model = keras.models.Model( - model.get_layer(name="image").input, model.get_layer( - name="dense3").output + model.get_layer(name="image").input, last_dense_layer.output ) + prediction_model.summary(line_length=110) inference_dataset = inference_generator