Update app.py

app.py

@@ -14,9 +14,9 @@ cfg = {
     "patch_size": 4,
     "in_channels": 3,
     "num_classes": 100,
-    "emb_dim": 768,
-    "num_heads": 12,
-    "depth": 12,
+    "emb_dim": 384,
+    "num_heads": 6,
+    "depth": 8,
     "mlp_ratio": 4.0,
     "drop": 0.1
 }
@@ -43,16 +43,33 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # ------------------------
 # Model definition
-class PatchEmbed(nn.Module):
-    def __init__(self, img_size=32, patch_size=4, in_chans=3, embed_dim=768):
+# ViT model implementation
+# --- Conv stem (replace PatchEmbed) ---
+class ConvPatchEmbed(nn.Module):
+    def __init__(self, in_chans=3, embed_dim=384):
         super().__init__()
-        self.patch_size = patch_size
-        self.n_patches = (img_size // patch_size) ** 2
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+        # Input 32x32 -> conv1: 32x32 -> conv2 stride2 -> 16x16 -> conv3 stride2 -> 8x8
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_chans, 64, kernel_size=3, stride=1, padding=1, bias=False),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(128),
+            nn.ReLU(inplace=True),
+
+            nn.Conv2d(128, embed_dim, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(embed_dim),
+            nn.ReLU(inplace=True),
+        )
+        # n_patches = (32/4)^2 = 8*8 = 64
+        self.n_patches = (32 // 4) ** 2
 
     def forward(self, x):
-        x = self.proj(x)
-        x = x.flatten(2).transpose(1,2)
+        # x: (B, C, H, W)
+        x = self.conv(x)                  # (B, E, H/4, W/4) -> H/4=8 for 32x32
+        x = x.flatten(2)                  # (B, E, N)
+        x = x.transpose(1, 2)             # (B, N, E)
         return x
 
 class MLP(nn.Module):
@@ -77,60 +94,75 @@ class Attention(nn.Module):
         self.num_heads = num_heads
         head_dim = dim // num_heads
         self.scale = head_dim ** -0.5
+
         self.qkv = nn.Linear(dim, dim*3, bias=qkv_bias)
         self.attn_drop = nn.Dropout(attn_drop)
         self.proj = nn.Linear(dim, dim)
         self.proj_drop = nn.Dropout(proj_drop)
+
     def forward(self, x):
         B, N, C = x.shape
-        qkv = self.qkv(x).reshape(B,N,3,self.num_heads,C//self.num_heads).permute(2,0,3,1,4)
-        q,k,v = qkv[0], qkv[1], qkv[2]
-        attn = (q @ k.transpose(-2,-1)) * self.scale
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2,0,3,1,4)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # each: (B, heads, N, head_dim)
+        attn = (q @ k.transpose(-2, -1)) * self.scale
         attn = attn.softmax(dim=-1)
         attn = self.attn_drop(attn)
-        x = (attn @ v).transpose(1,2).reshape(B,N,C)
+        x = (attn @ v).transpose(1,2).reshape(B, N, C)
         x = self.proj(x)
         x = self.proj_drop(x)
         return x
 
-class _StochasticDepth(nn.Module):
-    def __init__(self, p):
-        super().__init__()
-        self.p = p
-    def forward(self, x):
-        if not self.training or self.p==0.:
-            return x
-        keep = torch.rand(x.shape[0],1,1, device=x.device) >= self.p
-        return x * keep / (1 - self.p)
-
 class Block(nn.Module):
-    def __init__(self, dim, num_heads, mlp_ratio=4., drop=0., drop_path=0.):
+    def __init__(self, dim, num_heads, mlp_ratio=4., drop=0., attn_drop=0., drop_path=0.):
         super().__init__()
         self.norm1 = nn.LayerNorm(dim)
-        self.attn = Attention(dim, num_heads=num_heads, attn_drop=drop, proj_drop=drop)
-        self.drop_path = nn.Identity() if drop_path==0. else _StochasticDepth(drop_path)
+        self.attn = Attention(dim, num_heads=num_heads, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = nn.Identity() if drop_path == 0. else _StochasticDepth(drop_path)
         self.norm2 = nn.LayerNorm(dim)
         self.mlp = MLP(dim, int(dim*mlp_ratio), drop=drop)
+
     def forward(self, x):
         x = x + self.drop_path(self.attn(self.norm1(x)))
         x = x + self.drop_path(self.mlp(self.norm2(x)))
         return x
 
+# Simple implementation of stochastic depth
+class _StochasticDepth(nn.Module):
+    def __init__(self, p):
+        super().__init__()
+        self.p = p
+    def forward(self, x):
+        if not self.training or self.p == 0.:
+            return x
+        keep = torch.rand(x.shape[0], 1, 1, device=x.device) >= self.p
+        return x * keep / (1 - self.p)
+
 class ViT(nn.Module):
     def __init__(self, cfg):
         super().__init__()
-        self.patch_embed = PatchEmbed(cfg["image_size"], cfg["patch_size"], cfg["in_channels"], cfg["emb_dim"])
+        img_size, patch_size = cfg["image_size"], cfg["patch_size"]
+        # Use ConvPatchEmbed (hybrid) instead of linear patch conv with kernel=patch_size
+        self.patch_embed = ConvPatchEmbed(cfg["in_channels"], cfg["emb_dim"])
         n_patches = self.patch_embed.n_patches
+
         self.cls_token = nn.Parameter(torch.zeros(1,1,cfg["emb_dim"]))
-        self.pos_embed = nn.Parameter(torch.zeros(1, 1+n_patches, cfg["emb_dim"]))
+        self.pos_embed = nn.Parameter(torch.zeros(1, 1 + n_patches, cfg["emb_dim"]))
         self.pos_drop = nn.Dropout(p=cfg["drop"])
-        dpr = [x.item() for x in torch.linspace(0, 0.1, cfg["depth"])]
-        self.blocks = nn.ModuleList([Block(cfg["emb_dim"], cfg["num_heads"], cfg["mlp_ratio"], cfg["drop"], dpr[i]) for i in range(cfg["depth"])])
+
+        # transformer blocks
+        dpr = [x.item() for x in torch.linspace(0, cfg.get("drop_path", 0.2), cfg["depth"])]  # stochastic depth decay
+        self.blocks = nn.ModuleList([
+            Block(cfg["emb_dim"], num_heads=cfg["num_heads"], mlp_ratio=cfg["mlp_ratio"], drop=cfg["drop"], drop_path=dpr[i])
+            for i in range(cfg["depth"])
+        ])
         self.norm = nn.LayerNorm(cfg["emb_dim"])
         self.head = nn.Linear(cfg["emb_dim"], cfg["num_classes"])
-        nn.init.trunc_normal_(self.pos_embed,std=.02)
-        nn.init.trunc_normal_(self.cls_token,std=.02)
+
+        # init
+        nn.init.trunc_normal_(self.pos_embed, std=.02)
+        nn.init.trunc_normal_(self.cls_token, std=.02)
         self.apply(self._init_weights)
+
     def _init_weights(self, m):
         if isinstance(m, nn.Linear):
             nn.init.xavier_uniform_(m.weight)
@@ -139,17 +171,24 @@ class ViT(nn.Module):
         elif isinstance(m, nn.LayerNorm):
             nn.init.zeros_(m.bias)
             nn.init.ones_(m.weight)
-    def forward(self,x):
+        elif isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
+            if getattr(m, "bias", None) is not None:
+                nn.init.zeros_(m.bias)
+
+    def forward(self, x):
         B = x.shape[0]
-        x = self.patch_embed(x)
-        cls_tokens = self.cls_token.expand(B,-1,-1)
-        x = torch.cat((cls_tokens,x),dim=1)
+        x = self.patch_embed(x)             # (B, N, E)
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)   # (B, 1+N, E)
         x = x + self.pos_embed
         x = self.pos_drop(x)
+
         for blk in self.blocks:
             x = blk(x)
+
         x = self.norm(x)
-        cls = x[:,0]
+        cls = x[:, 0]
         out = self.head(cls)
         return out
 
@@ -180,7 +219,7 @@ iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(num_top_classes=1),
-    title="ViT CIFAR-100 Classifier",
+    title="Hybrid ViT+CNN CIFAR-100 Classifier",
     description="Upload a 32x32 image, and the model predicts the CIFAR-100 class."
 )