replace backbone with resnet18

models/base.py

@@ -1,5 +1,7 @@
 import tensorflow as tf
 from convnext.models.convnext_tf import ConvNeXt, ConvNeXtReverse, Block, ConvNeXtReverseLabels
+from models.resnet import resnet_18
+
 
 class Bottleneck(tf.keras.Model):
     
@@ -150,7 +152,8 @@ class AEEncoderNew(tf.keras.Model):
     def __init__(self, depths, dims, kernel_size=7, activation='tanh'):
         super(AEEncoderNew, self).__init__()
 
-        self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
+        # self.encoder = ConvNeXt(depths=depths, dims=dims, include_top=False, kernel_size=kernel_size)
+        self.encoder = resnet_18()
         self.act = tf.keras.layers.Activation(activation)
 
 
@@ -177,7 +180,8 @@ class DecoderNew(tf.keras.Model):
         self.is_label = is_label
 
         if not is_label:
-            self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
+            # self.decoder = ConvNeXtReverse(depths=depths, dims=dims, kernel_size=kernel_size)
+            self.decoder = resnet_18(reversed=True)
         else:
             self.decoder = ConvNeXtReverseLabels(depths=depths, dims=dims, nr_of_classes=7)