from unittest import skip import numpy as np import tensorflow as tf from tensorflow.keras.layers import (Conv2D, Conv2DTranspose, Dense, DepthwiseConv2D, GlobalAveragePooling2D, Layer, 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 if (not is_training) or (drop_prob == 0.): return inputs # Compute keep_prob keep_prob = 1.0 - drop_prob # Compute drop_connect tensor random_tensor = keep_prob shape = (tf.shape(inputs)[0],) + (1,) * \ (len(tf.shape(inputs)) - 1) random_tensor += tf.random.uniform(shape, dtype=inputs.dtype) binary_tensor = tf.floor(random_tensor) output = tf.math.divide(inputs, keep_prob) * binary_tensor return output class DropPath(tf.keras.layers.Layer): # borrowed from https://github.com/rishigami/Swin-Transformer-TF/blob/main/swintransformer/model.py def __init__(self, drop_prob=None): super().__init__() self.drop_prob = drop_prob def call(self, x, training): return drop_path(x, self.drop_prob, training) class Block(Layer): r""" ConvNeXt Block. There are two equivalent implementations: (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back We use (2) as we find it slightly faster in PyTorch Args: dim (int): Number of input channels. drop_path (float): Stochastic depth rate. Default: 0.0 layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. """ def __init__(self, dim, kernel_size=7, drop_path=0., layer_scale_init_value=1e-6, prefix=''): super().__init__() self.dwconv = DepthwiseConv2D( kernel_size=kernel_size, padding='same') # depthwise conv self.norm = LayerNormalization(epsilon=1e-6) # pointwise/1x1 convs, implemented with linear layers self.pwconv1 = Dense(4 * dim) self.act = tf.keras.activations.gelu self.pwconv2 = Dense(dim) self.drop_path = DropPath(drop_path) self.dim = dim self.layer_scale_init_value = layer_scale_init_value self.prefix = prefix def build(self, input_shape): self.gamma = tf.Variable( initial_value=self.layer_scale_init_value * tf.ones((self.dim)), trainable=True, name=f'{self.prefix}/gamma') self.built = True def call(self, x, training): input = 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, training=training) x = self.pwconv1(x, training=training) x = self.act(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, training=training) return x class ConvNeXt(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, kernel_size=7, 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.include_top = include_top 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)] # ) stem = [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')] # ) downsample_layer = [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], kernel_size=kernel_size, 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, training): for i in range(len(self.downsample_layers)): # x = self.downsample_layers[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, training): x = self.forward_features(x, training=training) return x 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], kernel_size=7, 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], kernel_size=kernel_size, 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, 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 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(len(depths) - 1): 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(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(nr_of_classes, activation='softmax') self.last = Dense(nr_of_classes - 1) def call(self, x, training): for i in reversed(range(len(self.downsample_layers))): x = self.downsample_layers[i](x, training=training) x = self.stages[i](x, training=training) x = self.last(x, training=training) return x model_urls = { "convnext_tiny_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_tiny_1k_224_ema.h5", "convnext_small_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_small_1k_224_ema.h5", "convnext_base_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_224.h5", "convnext_base_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_base_22k_1k_384.h5", "convnext_large_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_224.h5", "convnext_large_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_large_22k_1k_384.h5", "convnext_xlarge_224": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_224_ema.h5", "convnext_xlarge_384": "https://github.com/bamps53/convnext-tf/releases/download/v0.1/convnext_xlarge_22k_1k_384_ema.h5", } model_configs = dict( convnext_tiny=dict( depths=[3, 3, 9, 3], dims=[8, 16, 32, 512] ), convnext_small=dict( depths=[3, 3, 27, 3], dims=[96, 192, 384, 768] ), convnext_base=dict( depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024] ), convnext_large=dict( depths=[3, 3, 27, 3], dims=[192, 384, 768, 1536] ), convnext_xlarge=dict( depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048] ), )