added activations

francesco-vaselli · May 21, 2024 · b46c254 · b46c254
1 parent dbc326c
commit b46c254
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/configs/new_train_config.yaml b/configs/new_train_config.yaml
@@ -27,13 +27,15 @@ model_config:
       n_dense_layers: 2
       dense_size: 32
       output_shape: 6
+      output_activation: null
     rnn_config:
       n_rnn_layers: 5
       rnn_units: 512
       n_dense_layers: 6
       dense_size: 512
       output_shape: 6
       input_shape: [7, 1]
+      output_activation: null
     attn_config:
       input_shape: [7, 1]
       output_shape: 6
@@ -54,13 +56,15 @@ model_config:
       n_dense_layers: 1
       dense_size: 32
       output_shape: 1
+      output_activation: "sigmoid"
     rnn_config:
       n_rnn_layers: 3
       rnn_units: 416
       n_dense_layers: 6
       dense_size: 192
       output_shape: 1
       input_shape: [7, 1]
+      output_activation: "sigmoid"
     attn_config:
       input_shape: [7, 1]
       output_shape: 1
@@ -81,10 +85,12 @@ model_config:
       n_dense_layers: 4
       dense_size: 128
       output_shape: 3
+      output_activation: "softmax"
     rnn_config:
       n_rnn_layers: 4
       rnn_units: 512
       n_dense_layers: 2
       dense_size: 512
       output_shape: 3
-      input_shape: [7, 1]
+      input_shape: [7, 1]
+      output_activation: "softmax"
diff --git a/src/models/new_train_model.py b/src/models/new_train_model.py
@@ -137,7 +137,7 @@ def build_model(model_config, model_name):
             )
 
         # Add output layer
-        model.add(Dense(cnn_config["output_shape"]))
+        model.add(Dense(cnn_config["output_shape"], activation=cnn_config["output_activation"]))
         print("CNN model built:", "\n")
         model.summary()
 
@@ -161,7 +161,7 @@ def build_model(model_config, model_name):
             model.add(Dense(rnn_config["dense_size"], activation="relu"))
 
         # Add output layer
-        model.add(Dense(rnn_config["output_shape"]))
+        model.add(Dense(rnn_config["output_shape"], activation=rnn_config["output_activation"]))
         print("RNN model done")
         model.summary()
     elif model_name == "ar_rnn":