working weight learning

convnext/models/convnext_tf.py

@@ -2,7 +2,7 @@ from unittest import skip
 import numpy as np
 import tensorflow as tf
 from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
-                                     GlobalAveragePooling2D, Layer,
+                                     GlobalAveragePooling2D, Layer, Activation, Lambda,
                                      LayerNormalization, Flatten, Reshape, UpSampling2D)
 
 def drop_path(inputs, drop_prob, is_training):
@@ -257,12 +257,25 @@ class ConvNeXtReverse(tf.keras.Model):
 
         self.last = Dense(3, activation='tanh')
 
+        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
+        return x, weights
 
 
 class ConvNeXtReverseWithSkip(tf.keras.Model):