commit all

convnext/models/convnext_tf.py

@@ -29,7 +29,7 @@ class DropPath(tf.keras.layers.Layer):
         super().__init__()
         self.drop_prob = drop_prob
 
-    def call(self, x, training=None):
+    def call(self, x, training):
         return drop_path(x, self.drop_prob, training)
 
 
@@ -66,19 +66,19 @@ class Block(Layer):
             name=f'{self.prefix}/gamma')
         self.built = True
 
-    def call(self, x):
+    def call(self, x, training):
         input = x
-        x = self.dwconv(x)
+        x = self.dwconv(x, training=training)
         # x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
-        x = self.norm(x)
-        x = self.pwconv1(x)
+        x = self.norm(x, training=training)
+        x = self.pwconv1(x, training=training)
         x = self.act(x)
-        x = self.pwconv2(x)
+        x = self.pwconv2(x, training=training)
         if self.gamma is not None:
             x = self.gamma * x
         # x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
 
-        x = input + self.drop_path(x)
+        x = input + self.drop_path(x, training=training)
         return x
 
 
@@ -138,16 +138,16 @@ class ConvNeXt(tf.keras.Model):
             self.norm = None
             self.head = None
 
-    def forward_features(self, x):
+    def forward_features(self, x, training):
         for i in range(len(self.downsample_layers)):
             # x = self.downsample_layers[i](x)
-            x = self.downsample_layers[i][0](x)
-            x = self.downsample_layers[i][1](x)
-            x = self.stages[i](x)
+            x = self.downsample_layers[i][0](x, training=training)
+            x = self.downsample_layers[i][1](x, training=training)
+            x = self.stages[i](x, training=training)
         return x
 
-    def call(self, x):
-        x = self.forward_features(x)
+    def call(self, x, training):
+        x = self.forward_features(x, training=training)
         return x
 
 
@@ -257,11 +257,11 @@ class ConvNeXtReverse(tf.keras.Model):
 
         self.last = Dense(3, activation='tanh')
 
-    def call(self, x):
+    def call(self, x, training):
         for i in reversed(range(len(self.upsample_layers))):
-            x = self.upsample_layers[i](x)
-            x = self.stages[i](x)
-        x = self.last(x)
+            x = self.upsample_layers[i](x, training=training)
+            x = self.stages[i](x, training=training)
+        x = self.last(x, training=training)
         return x
 
 
@@ -366,13 +366,14 @@ class ConvNeXtReverseLabels(tf.keras.Model):
             self.stages.append(stage)
             cur += depths[i]
 
-        self.last = Dense(nr_of_classes, activation='softmax')
+        # self.last = Dense(nr_of_classes, activation='softmax')
+        self.last = Dense(nr_of_classes - 1)
 
-    def call(self, x):
+    def call(self, x, training):
         for i in reversed(range(len(self.downsample_layers))):
-            x = self.downsample_layers[i](x)
-            x = self.stages[i](x)
-        x = self.last(x)
+            x = self.downsample_layers[i](x, training=training)
+            x = self.stages[i](x, training=training)
+        x = self.last(x, training=training)
         return x