Update augvit_model.py

augvit_model.py

@@ -98,6 +98,33 @@ class Transformer(Layer):
             x = attn(x, training=training) + x + aug_attn(x, training=training)
             x = mlp(x, training=training) + x + augs(x, training=training)
         return x
+        
+@tf.keras.utils.register_keras_serializable()
+class AddPositionEmbs(tf.keras.layers.Layer):
+
+    def build(self, input_shape):
+        assert (
+            len(input_shape) == 3
+        ), f"Number of dimensions should be 3, got {len(input_shape)}"
+        self.pe = tf.Variable(
+            name="pos_embedding",
+            initial_value=tf.random_normal_initializer(stddev=0.06)(
+                shape=(1, input_shape[1], input_shape[2])
+            ),
+            dtype="float32",
+            trainable=True,
+        )
+
+    def call(self, inputs):
+        return inputs + tf.cast(self.pe, dtype=inputs.dtype)
+
+    def get_config(self):
+        config = super().get_config()
+        return config
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(**config)
 
 class AUGViT(Model):
     def __init__(self, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim,name='augvit',
@@ -117,11 +144,11 @@ class AUGViT(Model):
         self.patch_den=    nn.Dense(units=dim,name='patchden')
         
 
-       
+        self.pos_embedding = AddPositionEmbs(name="Transformer/posembed_input")
         self.cls_token = tf.Variable(initial_value=tf.random.normal([1, 1, dim]),name='cls',trainable=True)
         self.dropout = nn.Dropout(rate=emb_dropout,name='drop')
-        self.pos_embedding = tf.Variable(initial_value=tf.random_normal_initializer(stddev=0.06)(
-                shape=(1, num_patches + 1, dim)),name='pos_emb',trainable=True)
+        # self.pos_embedding = tf.Variable(initial_value=tf.random_normal_initializer(stddev=0.06)(
+        #         shape=(1, num_patches + 1, dim)),name='pos_emb',trainable=True)
         self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout=dropout,name='trans')
 
         self.pool = pool
@@ -142,7 +169,7 @@ class AUGViT(Model):
         )
         x = tf.concat([cls_tokens, x], axis=1)
         # print(x.shape,cls_tokens.shape )
-        x= x+self.pos_embedding
+        x= self.pos_embedding(x)
         
         # print(x.shape,pos.shape,self.pos_embedding.shape)
         x = self.dropout(x, training=training)