updated

vit.py

@@ -10,7 +10,7 @@ class ClassToken(layers.Layer):
         #initial values for the weight
         w_init = tf.random_normal_initializer()
         self.w = tf.Variable(
-            initial_value = w_init(shape=(1, 1, input_shape[-1]), dtype=tf.float32), 
+            initial_value = w_init(shape=(1, 1, input_shape[-1]), dtype=tf.float32),
             trainable = True
         )
 
@@ -22,7 +22,8 @@ class ClassToken(layers.Layer):
         cls = tf.broadcast_to(self.w, [batch_size, 1, hidden_dim])
         #change data type
         cls = tf.cast(cls, dtype=inputs.dtype)
-        return cls   
+        return cls
+
 
 def mlp(x, cf):
     x = layers.Dense(cf['mlp_dim'], activation='gelu')(x)
@@ -31,20 +32,20 @@ def mlp(x, cf):
     x = layers.Dropout(cf['dropout_rate'])(x)
     return x
 
-
 def transformer_encoder(x, cf):
     skip_1 = x
     x = layers.LayerNormalization()(x)
     x = layers.MultiHeadAttention(num_heads=cf['num_heads'], key_dim=cf['hidden_dim'])(x,x)
     x = layers.Add()([x, skip_1])
-    
+
     skip_2 = x
     x = layers.LayerNormalization()(x)
     x = mlp(x, cf)
     x = layers.Add()([x, skip_2])
-    
+
     return x
 
+
 def resnet_block(x, filters, strides=1):
     identity = x
 
@@ -63,13 +64,14 @@ def resnet_block(x, filters, strides=1):
     x = layers.Activation('relu')(x)
     return x
 
+
 def build_resnet(input_shape):
 
     x = layers.Conv2D(32, kernel_size=7, strides=2, padding='same')(input_shape)
     x = layers.BatchNormalization()(x)
     x = layers.Activation('relu')(x)
     x = layers.MaxPooling2D(pool_size=3, strides=2, padding='same')(x)
-    
+
     x = resnet_block(x, filters=32)
     x = resnet_block(x, filters=32)
 
@@ -78,13 +80,10 @@ def build_resnet(input_shape):
 
     x = resnet_block(x, filters=128, strides=2)
     x = resnet_block(x, filters=128)
-    
+
     x = resnet_block(x, filters=256, strides=2)
     x = resnet_block(x, filters=256)
-    
-    x = resnet_block(x, filters=512, strides=2)
-    x = resnet_block(x, filters=512)
-    
+
     return x
 
 
@@ -108,21 +107,20 @@ def CNN_ViT(hp):
     print(f"position embeding : {pos_embed.shape}")
     #Patch + Position Embedding
     embed = patch_embed + pos_embed
-    
+
     #Token
     token = ClassToken()(embed)
     x = layers.Concatenate(axis=1)([token, embed]) #(None, 257, 256)
-    
+
     #Transformer encoder
     for _ in range(hp['num_layers']):
         x = transformer_encoder(x, hp)
-    
-    
+
+
     x = layers.LayerNormalization()(x)
     x = x[:, 0, :]
     x = layers.Dense(hp['num_classes'], activation='softmax')(x)
-    
+
     model = Model(inputs, x)
-    
     return model