From e333dcba9c88cf993852340d7d3be4959db08522 Mon Sep 17 00:00:00 2001 From: Starbeamrainbowlabs Date: Fri, 16 Sep 2022 15:36:01 +0100 Subject: [PATCH] tweak projection head --- aimodel/src/lib/ai/model_rainfallwater_segmentation.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/aimodel/src/lib/ai/model_rainfallwater_segmentation.py b/aimodel/src/lib/ai/model_rainfallwater_segmentation.py index 127605c..fd70a1f 100644 --- a/aimodel/src/lib/ai/model_rainfallwater_segmentation.py +++ b/aimodel/src/lib/ai/model_rainfallwater_segmentation.py @@ -20,11 +20,9 @@ def model_rainfallwater_segmentation(metadata, feature_dim_in, shape_water_out, layer_next = do_convnext_inverse(layer_next, arch_name="convnext_i_tiny") - # TODO: An attention layer here instead of a dense layer, with a skip connection? + # TODO: An attention layer here instead of a dense layer, with a skip connection perhaps? layer_next = tf.keras.layers.Dense(32)(layer_next) - layer_next = tf.keras.layers.Conv2D(out_water_channels, 7, activation="softmax", padding="same")(layer_next) - - # TODO: Implement projection head here + layer_next = tf.keras.layers.Conv2D(out_water_channels, kernel_size=1, activation="softmax", padding="same")(layer_next) model = tf.keras.Model( inputs = layer_input,