From f2e2ca14840d1e65eea115c028f14a231fb45f5d Mon Sep 17 00:00:00 2001 From: Starbeamrainbowlabs Date: Mon, 24 Oct 2022 20:52:31 +0100 Subject: [PATCH] model_contrastive: make water encoder significantly shallower --- aimodel/src/lib/ai/components/convnext.py | 1 + aimodel/src/lib/ai/model_rainfallwater_contrastive.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/aimodel/src/lib/ai/components/convnext.py b/aimodel/src/lib/ai/components/convnext.py index 0202ae7..b49fb38 100644 --- a/aimodel/src/lib/ai/components/convnext.py +++ b/aimodel/src/lib/ai/components/convnext.py @@ -11,6 +11,7 @@ kernel_initial = tf.keras.initializers.TruncatedNormal(stddev=0.2) bias_initial = tf.keras.initializers.Constant(value=0) depths_dims = dict( + convnext_shallow= (dict(depths=[1, 3, 1], dims=[66, 132, 192])), convnext_xtiny = (dict(depths=[3, 3, 6, 3], dims=[66, 132, 264, 528])), # architectures from: https://github.com/facebookresearch/ConvNeXt # A ConvNet for the 2020s: https://arxiv.org/abs/2201.03545 diff --git a/aimodel/src/lib/ai/model_rainfallwater_contrastive.py b/aimodel/src/lib/ai/model_rainfallwater_contrastive.py index 80eed22..cf970ed 100644 --- a/aimodel/src/lib/ai/model_rainfallwater_contrastive.py +++ b/aimodel/src/lib/ai/model_rainfallwater_contrastive.py @@ -44,7 +44,7 @@ def model_rainfallwater_contrastive(metadata, shape_water, batch_size=64, featur input_height=water_height, input_channels=water_channels, feature_dim=feature_dim, - arch_name="convnext_xtiny", + arch_name="convnext_shallow", summary_file=summary_file )(input_water)