new models

gcvit/__init__.py

@@ -1,2 +1,2 @@
-from .models import GCViT, GCViTTiny, GCViTSmall, GCViTBase
+from .models import GCViT, GCViTXXTiny, GCViTXTiny, GCViTTiny, GCViTSmall, GCViTBase, GCViTLarge
 from .version import __version__

gcvit/models/__init__.py

@@ -1 +1 @@
-from .gcvit import GCViT, GCViTTiny, GCViTSmall, GCViTBase
+from .gcvit import GCViT, GCViTXXTiny, GCViTXTiny, GCViTTiny, GCViTSmall, GCViTBase, GCViTLarge

gcvit/models/gcvit.py

@@ -1,16 +1,29 @@
 import numpy as np
 import tensorflow as tf
 
-from ..layers import PatchEmbed, GCViTLayer, Identity
-
+from ..layers import Stem, GCViTLevel, Identity
 
+                  
 BASE_URL = 'https://github.com/awsaf49/gcvit-tf/releases/download'
-TAG = 'v1.0.4'
+TAG = 'v1.1.1'
 NAME2CONFIG = {
+    'gcvit_xxtiny': {'window_size': (7, 7, 14, 7),
+                    'dim': 64,
+                    'depths': (2, 2, 6, 2),
+                    'num_heads': (2, 4, 8, 16),
+                    'mlp_ratio': 3.,
+                    'path_drop': 0.2},
+    'gcvit_xtiny': {'window_size': (7, 7, 14, 7),
+                    'dim': 64,
+                    'depths': (3, 4, 6, 5),
+                    'num_heads': (2, 4, 8, 16),
+                    'mlp_ratio': 3.,
+                    'path_drop': 0.2},
     'gcvit_tiny': {'window_size': (7, 7, 14, 7),
                     'dim': 64,
                     'depths': (3, 4, 19, 5),
-                    'num_heads': (2, 4, 8, 16), 
+                    'num_heads': (2, 4, 8, 16),
+                    'mlp_ratio': 3.,
                     'path_drop': 0.2,},
     'gcvit_small': {'window_size': (7, 7, 14, 7), 
                      'dim': 96, 
@@ -26,6 +39,13 @@ NAME2CONFIG = {
                     'mlp_ratio': 2.,
                     'path_drop': 0.5,
                     'layer_scale': 1e-5,},
+    'gcvit_large': {'window_size': (7, 7, 14, 7),
+                    'dim':192, 
+                    'depths': (3, 4, 19, 5),
+                    'num_heads': (6, 12, 24, 48),
+                    'mlp_ratio': 2.,
+                    'path_drop': 0.5,
+                    'layer_scale': 1e-5,},
     }
 
 @tf.keras.utils.register_keras_serializable(package='gcvit')
@@ -50,14 +70,14 @@ class GCViT(tf.keras.Model):
         self.num_classes = num_classes
         self.head_act = head_act
 
-        self.patch_embed = PatchEmbed(dim=dim, name='patch_embed')
+        self.patch_embed = Stem(dim=dim, name='patch_embed')
         self.pos_drop = tf.keras.layers.Dropout(drop_rate, name='pos_drop')
         path_drops = np.linspace(0., path_drop, sum(depths))
         keep_dims = [(False, False, False),(False, False),(True,),(True,),]
         self.levels = []
         for i in range(len(depths)):
             path_drop = path_drops[sum(depths[:i]):sum(depths[:i + 1])].tolist()
-            level = GCViTLayer(depth=depths[i], num_heads=num_heads[i], window_size=window_size[i], keep_dims=keep_dims[i],
+            level = GCViTLevel(depth=depths[i], num_heads=num_heads[i], window_size=window_size[i], keep_dims=keep_dims[i],
                     downsample=(i < len(depths) - 1), mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, 
                     drop=drop_rate, attn_drop=attn_drop, path_drop=path_drop, layer_scale=layer_scale, resize_query=resize_query,
                     name=f'levels/{i}')
@@ -71,16 +91,14 @@ class GCViT(tf.keras.Model):
             self.pool = Identity(name='pool')
         else:
             raise ValueError(f'Expecting pooling to be one of None/avg/max. Found: {global_pool}')
-        self.head = [tf.keras.layers.Dense(num_classes, name='head/fc'),
-                     tf.keras.layers.Activation(head_act, name='head/act')]
+        self.head = tf.keras.layers.Dense(num_classes, name='head', activation=head_act)
 
-    def reset_classifier(self, num_classes, head_act, global_pool=None):
+    def reset_classifier(self, num_classes, head_act, global_pool=None, in_channels=3):
         self.num_classes = num_classes
         if global_pool is not None:
             self.global_pool = global_pool
-        self.head[0] = tf.keras.layers.Dense(num_classes, name='head/fc') if num_classes else Identity(name='head/fc')
-        self.head[1] = tf.keras.layers.Activation(head_act, name='head/act') if head_act else Identity(name='head/act')
-        super().build((1, 224, 224, 3))
+        self.head = tf.keras.layers.Dense(num_classes, name='head', activation=head_act) if num_classes else Identity(name='head')
+        super().build((1, 224, 224, in_channels)) # for head we only need info from the input channel
         
     def forward_features(self, inputs):
         x = self.patch_embed(inputs)
@@ -96,8 +114,7 @@ class GCViT(tf.keras.Model):
         if self.global_pool in ['avg', 'max']:
             x = self.pool(x)
         if not pre_logits:
-            for layer in self.head:
-                x = layer(x)
+            x = self.head(x)
         return x 
 
     def call(self, inputs, **kwargs):
@@ -110,35 +127,71 @@ class GCViT(tf.keras.Model):
         x = tf.keras.Input(shape=input_shape)
         return tf.keras.Model(inputs=[x], outputs=self.call(x), name=self.name)
 
+    def summary(self, input_shape=(224, 224, 3)):
+        return self.build_graph(input_shape).summary()
+
 # load standard models
-def GCViTTiny(pretrain=False, **kwargs):
+def GCViTXXTiny(input_shape=(224, 224, 3), pretrain=False, resize_query=False, **kwargs):
+    name = 'gcvit_xxtiny'
+    config = NAME2CONFIG[name]
+    ckpt_link = '{}/{}/{}_weights.h5'.format(BASE_URL, TAG, name)
+    model = GCViT(name=name, resize_query=resize_query, **config,  **kwargs)
+    model(tf.random.uniform(shape=input_shape)[tf.newaxis,])
+    if pretrain:
+        ckpt_path = tf.keras.utils.get_file('{}_weights.h5'.format(name), ckpt_link)
+        model.load_weights(ckpt_path)
+    return model
+
+def GCViTXTiny(input_shape=(224, 224, 3), pretrain=False, resize_query=False, **kwargs):
+    name = 'gcvit_xtiny'
+    config = NAME2CONFIG[name]
+    ckpt_link = '{}/{}/{}_weights.h5'.format(BASE_URL, TAG, name)
+    model = GCViT(name=name, resize_query=resize_query, **config,  **kwargs)
+    model(tf.random.uniform(shape=input_shape)[tf.newaxis,])
+    if pretrain:
+        ckpt_path = tf.keras.utils.get_file('{}_weights.h5'.format(name), ckpt_link)
+        model.load_weights(ckpt_path)
+    return model
+
+def GCViTTiny(input_shape=(224, 224, 3), pretrain=False, resize_query=False, **kwargs):
     name = 'gcvit_tiny'
     config = NAME2CONFIG[name]
     ckpt_link = '{}/{}/{}_weights.h5'.format(BASE_URL, TAG, name)
-    model = GCViT(name=name, **config, **kwargs)
-    model(tf.random.uniform(shape=(1, 224, 224, 3)))
+    model = GCViT(name=name, resize_query=resize_query, **config,  **kwargs)
+    model(tf.random.uniform(shape=input_shape)[tf.newaxis,])
     if pretrain:
         ckpt_path = tf.keras.utils.get_file('{}_weights.h5'.format(name), ckpt_link)
         model.load_weights(ckpt_path)
     return model
 
-def GCViTSmall(pretrain=False, **kwargs):
+def GCViTSmall(input_shape=(224, 224, 3), pretrain=False, resize_query=False, **kwargs):
     name = 'gcvit_small'
     config = NAME2CONFIG[name]
     ckpt_link = '{}/{}/{}_weights.h5'.format(BASE_URL, TAG, name)
-    model = GCViT(name=name, **config, **kwargs)
-    model(tf.random.uniform(shape=(1, 224, 224, 3)))
+    model = GCViT(name=name, resize_query=resize_query, **config, **kwargs)
+    model(tf.random.uniform(shape=input_shape)[tf.newaxis,])
     if pretrain:
         ckpt_path = tf.keras.utils.get_file('{}_weights.h5'.format(name), ckpt_link)
         model.load_weights(ckpt_path)
     return model
 
-def GCViTBase(pretrain=False, **kwargs):
+def GCViTBase(input_shape=(224, 224, 3), pretrain=False, resize_query=False, **kwargs):
     name = 'gcvit_base'
     config = NAME2CONFIG[name]
     ckpt_link = '{}/{}/{}_weights.h5'.format(BASE_URL, TAG, name)
-    model = GCViT(name=name, **config, **kwargs)
-    model(tf.random.uniform(shape=(1, 224, 224, 3)))
+    model = GCViT(name=name, resize_query=resize_query, **config, **kwargs)
+    model(tf.random.uniform(shape=input_shape)[tf.newaxis,])
+    if pretrain:
+        ckpt_path = tf.keras.utils.get_file('{}_weights.h5'.format(name), ckpt_link)
+        model.load_weights(ckpt_path)
+    return model
+
+def GCViTLarge(input_shape=(224, 224, 3), pretrain=False, resize_query=False, **kwargs):
+    name = 'gcvit_large'
+    config = NAME2CONFIG[name]
+    ckpt_link = '{}/{}/{}_weights.h5'.format(BASE_URL, TAG, name)
+    model = GCViT(name=name, resize_query=resize_query, **config, **kwargs)
+    model(tf.random.uniform(shape=input_shape)[tf.newaxis,])
     if pretrain:
         ckpt_path = tf.keras.utils.get_file('{}_weights.h5'.format(name), ckpt_link)
         model.load_weights(ckpt_path)