Create vit.py

models/vit.py

@@ -0,0 +1,110 @@
+import tensorflow as tf
+from tensorflow import keras
+from tensorflow.keras import layers
+
+# Model hyperparameters
+num_classes = 3
+input_shape = (256, 256, 3)
+image_size = 256
+patch_size = 16
+num_patches = (image_size // patch_size) ** 2
+projection_dim = 64
+num_heads = 8  # Increased from 4 → 8
+transformer_units = [projection_dim * 2, projection_dim]
+transformer_layers = 12  # Increased from 8 → 12
+mlp_head_units = [2048, 1024]
+dropout_rate = 0.1
+
+# Data Augmentation with stronger transformations
+data_augmentation = keras.Sequential(
+    [
+        layers.Normalization(),
+        layers.Resizing(image_size, image_size),
+        layers.RandomFlip("horizontal"),
+        layers.RandomRotation(factor=0.1),
+        layers.RandomZoom(0.2, 0.2),
+    ],
+    name="data_augmentation",
+)
+
+# Patch Creation Layer
+class Patches(layers.Layer):
+    def __init__(self, patch_size):
+        super().__init__()
+        self.patch_size = patch_size
+
+    def call(self, images):
+        batch_size = tf.shape(images)[0]
+        patches = tf.image.extract_patches(
+            images=images,
+            sizes=[1, self.patch_size, self.patch_size, 1],
+            strides=[1, self.patch_size, self.patch_size, 1],
+            rates=[1, 1, 1, 1],
+            padding="VALID",
+        )
+        patch_dims = patches.shape[-1]
+        return tf.reshape(patches, [batch_size, -1, patch_dims])
+
+# Patch Encoding with Learnable [CLS] Token
+class PatchEncoder(layers.Layer):
+    def __init__(self, num_patches, projection_dim):
+        super().__init__()
+        self.projection = layers.Dense(projection_dim)
+        self.cls_token = self.add_weight(shape=(1, 1, projection_dim), initializer="zeros", trainable=True)
+        self.position_embedding = layers.Embedding(input_dim=num_patches + 1, output_dim=projection_dim)
+
+    def call(self, patch):
+        batch_size = tf.shape(patch)[0]
+        cls_token = tf.broadcast_to(self.cls_token, [batch_size, 1, projection_dim])
+        patch = self.projection(patch)
+        patch = tf.concat([cls_token, patch], axis=1)  # Add CLS token
+        positions = tf.range(start=0, limit=num_patches + 1, delta=1)
+        return patch + self.position_embedding(positions)
+
+# MLP Block
+def mlp(x, hidden_units, dropout_rate):
+    for units in hidden_units:
+        x = layers.Dense(units, activation=tf.nn.gelu)(x)
+        x = layers.Dropout(dropout_rate)(x)
+        x = layers.LayerNormalization()(x)  # Added LayerNorm
+    return x
+
+# Vision Transformer Model
+def create_vit_classifier():
+    inputs = layers.Input(shape=input_shape)
+    augmented = data_augmentation(inputs)
+    patches = Patches(patch_size)(augmented)
+    encoded_patches = PatchEncoder(num_patches, projection_dim)(patches)
+
+    for _ in range(transformer_layers):
+        x1 = layers.LayerNormalization()(encoded_patches)
+        attention_output = layers.MultiHeadAttention(num_heads=num_heads, key_dim=projection_dim, dropout=dropout_rate)(x1, x1)
+        x2 = layers.Add()([attention_output, encoded_patches])
+        x3 = layers.LayerNormalization()(x2)
+        x3 = mlp(x3, transformer_units, dropout_rate)
+        encoded_patches = layers.Add()([x3, x2])
+
+    representation = layers.LayerNormalization()(encoded_patches)
+    cls_output = representation[:, 0, :]  # Extract CLS token representation
+    features = mlp(cls_output, mlp_head_units, dropout_rate)
+    outputs = layers.Dense(num_classes, activation="softmax")(features)
+
+    return keras.Model(inputs, outputs)
+
+# Cosine Decay Learning Rate Scheduler
+def cosine_decay_schedule(initial_lr, total_steps):
+    return keras.optimizers.schedules.CosineDecay(initial_learning_rate=initial_lr, decay_steps=total_steps)
+
+# Optimizer with Weight Decay & Cosine Decay
+learning_rate = 3e-4
+weight_decay = 0.03
+num_epochs = 50
+
+optimizer = keras.optimizers.AdamW(learning_rate=cosine_decay_schedule(learning_rate, num_epochs * 100), weight_decay=weight_decay)
+
+# Compile Model
+vit_classifier = create_vit_classifier()
+vit_classifier.compile(optimizer=optimizer, loss="sparse_categorical_crossentropy", metrics=["accuracy"])
+
+# Print Model Summary
+vit_classifier.summary()