lot of things

2022-08-06 12:27:24 +00:00
parent cd498e79f3
commit ee8f1cb68c
28 changed files with 4595 additions and 175 deletions
@@ -1,9 +1,9 @@
+from unittest import skip
 import numpy as np
 import tensorflow as tf
-from tensorflow.keras.layers import (Conv2D, Dense, DepthwiseConv2D,
+from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D,
                                     GlobalAveragePooling2D, Layer,
-                                     LayerNormalization)
-
+                                     LayerNormalization, Flatten, Reshape, UpSampling2D)

 def drop_path(inputs, drop_prob, is_training):
    # borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py
@@ -109,6 +109,299 @@ class ConvNeXt(tf.keras.Model):
            LayerNormalization(epsilon=1e-6)]
        )
        self.downsample_layers.append(stem)
+        for i in range(len(depths) - 1):
+            downsample_layer = tf.keras.Sequential([
+                LayerNormalization(epsilon=1e-6),
+                Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(len(depths)):
+            stage = tf.keras.Sequential(
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        if self.include_top:
+            self.avg = GlobalAveragePooling2D()
+            self.norm = LayerNormalization(epsilon=1e-6)  # final norm layer
+            self.head = Dense(num_classes)
+        else:
+            self.avg = None
+            self.norm = None
+            self.head = None
+
+    def forward_features(self, x):
+        for i in range(len(self.downsample_layers)):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+        return x
+
+    def call(self, x):
+        x = self.forward_features(x)
+        if self.include_top:
+            x = self.avg(x)
+            x = self.norm(x)
+            x = self.head(x)
+        return x
+
+
+class ConvNeXtWithSkip(tf.keras.Model):
+    r""" ConvNeXt
+        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        include_top (bool): whether to add head or just use it as feature extractor. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(self, 
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], 
+                 drop_path_rate=0., layer_scale_init_value=1e-6, 
+                 ):
+        super().__init__()
+        self.downsample_layers = []  # stem and 3 intermediate downsampling conv layers
+        stem = tf.keras.Sequential([
+            Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
+            LayerNormalization(epsilon=1e-6)]
+        )
+        self.downsample_layers.append(stem)
+        for i in range(len(depths) - 1):
+            downsample_layer = tf.keras.Sequential([
+                LayerNormalization(epsilon=1e-6),
+                Conv2D(dims[i+1], kernel_size=2, strides=2, padding='same')]
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(len(depths)):
+            stage = tf.keras.Sequential(
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+
+    def call(self, x):
+        skip_features = []
+        for i in range(len(self.downsample_layers)):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+            skip_features.append(x)
+        return x, skip_features
+
+
+class ConvNeXtReverse(tf.keras.Model):
+    r""" ConvNeXt
+        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        include_top (bool): whether to add head or just use it as feature extractor. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(self, 
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
+                 ):
+        super().__init__()
+        self.upsample_layers = []  # stem and 3 intermediate downsampling conv layers
+        stem = tf.keras.Sequential([
+            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
+            # UpSampling2D((4, 4), interpolation='nearest'),
+            # Conv2D(dims[0], kernel_size=4, strides=1, padding='same'),
+            # Dense(dims[0]),
+            LayerNormalization(epsilon=1e-6)]
+        )
+        self.upsample_layers.append(stem)
+        for i in range(len(depths) - 1):
+            downsample_layer = tf.keras.Sequential([
+                LayerNormalization(epsilon=1e-6),
+                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')
+                # UpSampling2D((2,2), interpolation='nearest'),
+                # Conv2D(dims[i+1], kernel_size=2, strides=1, padding='same'),
+                # Dense(dims[i+1]),
+            ])
+            self.upsample_layers.append(downsample_layer)
+
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(len(depths)):
+            stage = tf.keras.Sequential(
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        self.last = Dense(3, activation='tanh')
+
+    def call(self, x):
+        for i in reversed(range(len(self.upsample_layers))):
+            x = self.upsample_layers[i](x)
+            x = self.stages[i](x)
+        x = self.last(x)
+        return x
+
+
+class ConvNeXtReverseWithSkip(tf.keras.Model):
+    r""" ConvNeXt
+        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        include_top (bool): whether to add head or just use it as feature extractor. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(self, 
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
+                 ):
+        super().__init__()
+        assert len(depths) == len(dims)
+        self.upsample_layers = []  # stem and 3 intermediate downsampling conv layers
+        stem = tf.keras.Sequential([
+            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
+            LayerNormalization(epsilon=1e-6)]
+        )
+        self.upsample_layers.append(stem)
+        for i in range(len(depths) - 1):
+            downsample_layer = tf.keras.Sequential([
+                LayerNormalization(epsilon=1e-6),
+                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
+            )
+            self.upsample_layers.append(downsample_layer)
+
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(len(depths)):
+            stage = tf.keras.Sequential(
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        self.last = Dense(3, activation='tanh')
+
+    def call(self, x):
+        x, skip_features = x
+        for i in reversed(range(len(self.upsample_layers))):
+            x = tf.concat((x, skip_features[i]), axis=-1)
+            x = self.upsample_layers[i](x)
+            x = self.stages[i](x)
+        x = self.last(x)
+        return x
+
+
+class ConvNeXtReverseLabels(tf.keras.Model):
+    r""" ConvNeXt
+        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        include_top (bool): whether to add head or just use it as feature extractor. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(self, 
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
+                 nr_of_classes=7):
+        super().__init__()
+        self.downsample_layers = []  # stem and 3 intermediate downsampling conv layers
+        stem = tf.keras.Sequential([
+            Conv2DTranspose(dims[0], kernel_size=4, strides=4, padding='same'),
+            LayerNormalization(epsilon=1e-6)]
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = tf.keras.Sequential([
+                LayerNormalization(epsilon=1e-6),
+                Conv2DTranspose(dims[i+1], kernel_size=2, strides=2, padding='same')]
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = []  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x for x in np.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(4):
+            stage = tf.keras.Sequential(
+                [Block(dim=dims[i], drop_path=dp_rates[cur + j],
+                       layer_scale_init_value=layer_scale_init_value, prefix=f'block{i}') for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        self.last = Dense(nr_of_classes, activation='softmax')
+
+    def call(self, x):
+        for i in reversed(range(4)):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+        x = self.last(x)
+        return x
+
+
+class ConvNeXtDiscriminator(tf.keras.Model):
+    r""" ConvNeXt
+        A Tensorflow keras impl of : `A ConvNet for the 2020s`  -
+          https://arxiv.org/pdf/2201.03545.pdf
+
+    Args:
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        include_top (bool): whether to add head or just use it as feature extractor. Default: True
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(self, num_classes=1000,
+                 depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], include_top=True,
+                 drop_path_rate=0., layer_scale_init_value=1e-6, head_init_scale=1.,
+                 ):
+        super().__init__()
+        self.downsample_layers = []  # stem and 3 intermediate downsampling conv layers
+        stem = tf.keras.Sequential([
+            Conv2D(dims[0], kernel_size=4, strides=4, padding='same'),
+            LayerNormalization(epsilon=1e-6)]
+        )
+        self.downsample_layers.append(stem)
        for i in range(3):
            downsample_layer = tf.keras.Sequential([
                LayerNormalization(epsilon=1e-6),
@@ -127,28 +420,21 @@ class ConvNeXt(tf.keras.Model):
            self.stages.append(stage)
            cur += depths[i]

-        if self.include_top:
-            self.avg = GlobalAveragePooling2D()
-            self.norm = LayerNormalization(epsilon=1e-6)  # final norm layer
-            self.head = Dense(num_classes)
-        else:
-            self.avg = None
-            self.norm = None
-            self.head = None
+        self.flatten = Flatten()
+        self.last = Dense(1)

-    def forward_features(self, x):
+
+    def call(self, x):
+        x = tf.transpose(x, [0,2,3,1])
        for i in range(4):
            x = self.downsample_layers[i](x)
            x = self.stages[i](x)
+        x = self.flatten(x)
+        x = self.last(x)
        return x

-    def call(self, x):
-        x = self.forward_features(x)
-        if self.include_top:
-            x = self.avg(x)
-            x = self.norm(x)
-            x = self.head(x)
-        return x
+
+


 model_urls = {