From 8c0efd10cd42023ab81a680ee278d7812ff80372 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Wed, 14 Sep 2022 16:53:55 +0000 Subject: [PATCH] fix curvature model --- convnext/models/convnext_tf.py | 6 +++--- models/ae_curv_model.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/convnext/models/convnext_tf.py b/convnext/models/convnext_tf.py index 67c8d8b..0387db0 100644 --- a/convnext/models/convnext_tf.py +++ b/convnext/models/convnext_tf.py @@ -348,7 +348,7 @@ class ConvNeXtReverseLabels(tf.keras.Model): LayerNormalization(epsilon=1e-6)] ) self.downsample_layers.append(stem) - for i in range(3): + for i in range(len(depths) - 1): downsample_layer = tf.keras.Sequential([ LayerNormalization(epsilon=1e-6), Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')] @@ -358,7 +358,7 @@ class ConvNeXtReverseLabels(tf.keras.Model): self.stages = [] # 4 feature resolution stages, each consisting of multiple residual blocks dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))] cur = 0 - for i in range(4): + for i in range(len(depths)): stage = tf.keras.Sequential( [Block(dim=dims[i], drop_path=dp_rates[cur + j], layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])] @@ -369,7 +369,7 @@ class ConvNeXtReverseLabels(tf.keras.Model): self.last = Dense(nr_of_classes, activation='softmax') def call(self, x): - for i in reversed(range(4)): + for i in reversed(len(depths)): x = self.downsample_layers[i](x) x = self.stages[i](x) x = self.last(x) diff --git a/models/ae_curv_model.py b/models/ae_curv_model.py index ad0d6f1..242b4de 100644 --- a/models/ae_curv_model.py +++ b/models/ae_curv_model.py @@ -74,7 +74,7 @@ class AutoEncoder(tf.keras.Model): # x = x * uv_mask[:,:, tf.newaxis] # x = x * uv_mask[:,tf.newaxis,:,:] - x = x * uv_mask[:,:,:,tf.newaxis] + # x = x * uv_mask[:,:,:,tf.newaxis] return x