apply for body and curv

models/resnet.py

@@ -116,9 +116,10 @@ def make_bottleneck_layer(filter_num, blocks, stride=1):
 
 
 class ResNetTypeI(tf.keras.Model):
-    def __init__(self, layer_params, reversed=False):
+    def __init__(self, layer_params, reversed=False, labels=None):
         super(ResNetTypeI, self).__init__()
         self.reversed = reversed
+        self.labels = labels
 
         if not reversed:
             self.conv1 = tf.keras.layers.Conv2D(filters=64,
@@ -130,7 +131,7 @@ class ResNetTypeI(tf.keras.Model):
                                                 strides=2,
                                                 padding="same")
         else:
-            self.conv1 = tf.keras.layers.Conv2DTranspose(filters=3,
+            self.conv1 = tf.keras.layers.Conv2DTranspose(filters=3 if labels is None else labels,
                                                 kernel_size=(7, 7),
                                                 strides=2,
                                                 padding="same")
@@ -177,7 +178,10 @@ class ResNetTypeI(tf.keras.Model):
             x = self.bn1(x, training=training)
             # x = tf.nn.relu(x)
             x = self.conv1(x)
-            x = tf.nn.tanh(x)
+            if self.labels is None:
+                x = tf.nn.tanh(x)
+            # else:
+            #     x = tf.nn.softmax(x)
 
         return x
 
@@ -224,8 +228,8 @@ class ResNetTypeII(tf.keras.Model):
         return output
 
 
-def resnet_18(reversed=False):
-    return ResNetTypeI(layer_params=[2, 2, 2, 2], reversed=reversed)
+def resnet_18(reversed=False, labels=None):
+    return ResNetTypeI(layer_params=[2, 2, 2, 2], reversed=reversed, labels=labels)
 
 
 def resnet_34():