convolutional weights model

convnext/models/convnext_tf.py

@@ -222,7 +222,7 @@ class ConvNeXtReverse(tf.keras.Model):
 
     def __init__(self, 
                  depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], kernel_size=7,
-                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1., learn_weights=False,
                  ):
         super().__init__()
         self.upsample_layers = []  # stem and 3 intermediate downsampling conv layers
@@ -255,27 +255,20 @@ class ConvNeXtReverse(tf.keras.Model):
             self.stages.append(stage)
             cur += depths[i]
 
-        self.last = Dense(3, activation='tanh')
+        if not learn_weights:
+            self.last = Dense(3, activation='tanh')
+        else:
+            self.last = Activation('softmax')
+            # self.last = Dense(24, activation='softmax')
 
-        self.weights_model = tf.keras.Sequential([Block(dim=dims[-1]*2, kernel_size=kernel_size),
-                                                  Block(dim=dims[-1]*2, kernel_size=kernel_size),
-                                                  Block(dim=dims[-1]*2, kernel_size=kernel_size),
-                                                  Flatten(),
-                                                  Dense(128 * 24),
-                                                  Reshape((24, 128)),
-                                                  Dense(14475),
-                                                  Lambda(lambda x: tf.transpose(x, (0, 2, 1))),
-                                                  Lambda(lambda x: tf.nn.softmax(x, axis=-1))])
 
     def call(self, x, training):
-        weights = self.weights_model(x, training=training)
-        
         for i in reversed(range(len(self.upsample_layers))):
             x = self.upsample_layers[i](x, training=training)
             x = self.stages[i](x, training=training)
 
         x = self.last(x, training=training)
-        return x, weights
+        return x
 
 
 class ConvNeXtReverseWithSkip(tf.keras.Model):