Update app.py

app.py

@@ -13,9 +13,9 @@ cfg = {
     "patch_size": 4,
     "in_channels": 3,
     "num_classes": 100,
-    "emb_dim": 384,
+    "emb_dim": 192,
     "num_heads": 6,
-    "depth": 8,
+    "depth": 6,
     "mlp_ratio": 4.0,
     "drop": 0.1
 }
@@ -40,32 +40,32 @@ classes = [
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ViT model implementation
 class ConvPatchEmbed(nn.Module):
-    def __init__(self, in_chans=3, embed_dim=384):
+    def __init__(self, img_size=32, in_chans=3, embed_dim=192):
         super().__init__()
-        # Input 32x32 -> conv1: 32x32 -> conv2 stride2 -> 16x16 -> conv3 stride2 -> 8x8
-        self.conv = nn.Sequential(
+        # 32x32 -> 32x32 -> 16x16 -> 16x16
+        self.proj = 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.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=False),  # 32 -> 16
             nn.BatchNorm2d(128),
             nn.ReLU(inplace=True),
 
-            nn.Conv2d(128, embed_dim, kernel_size=3, stride=2, padding=1, bias=False),
+            nn.Conv2d(128, embed_dim, kernel_size=3, stride=1, padding=1, bias=False),  # stays 16x16
             nn.BatchNorm2d(embed_dim),
             nn.ReLU(inplace=True),
         )
-        # n_patches = (32/4)^2 = 8*8 = 64
-        self.n_patches = (32 // 4) ** 2
+
+        grid_size = (img_size // 2, img_size // 2)  # (16,16)
+        self.grid_size = grid_size
+        self.num_patches = grid_size[0] * grid_size[1]
 
     def forward(self, x):
-        # x: (B, C, H, W)
-        x = self.conv(x)                 
-        x = x.flatten(2)                 
-        x = x.transpose(1, 2)             
+        x = self.proj(x)         # (B, E, H=16, W=16)
+        B, C, H, W = x.shape
+        x = x.flatten(2).transpose(1, 2)  # (B, N=H*W, E)
         return x
 
 class MLP(nn.Module):
@@ -76,6 +76,7 @@ class MLP(nn.Module):
         self.act = nn.GELU()
         self.fc2 = nn.Linear(hidden_features, in_features)
         self.drop = nn.Dropout(drop)
+
     def forward(self, x):
         x = self.fc1(x)
         x = self.act(x)
@@ -85,70 +86,87 @@ class MLP(nn.Module):
         return x
 
 class Attention(nn.Module):
-    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
+    def __init__(self, dim, num_heads=8, qkv_bias=True, attn_drop=0., proj_drop=0.):
         super().__init__()
         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.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]   # each: (B, heads, N, head_dim)
+        # (B, N, 3C) -> (3, B, heads, N, head_dim)
+        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
         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
 
+# Simple Stochastic Depth
+class StochasticDepth(nn.Module):
+    def __init__(self, p):
+        super().__init__()
+        self.p = float(p)
+
+    def forward(self, x):
+        if not self.training or self.p == 0.0:
+            return x
+        keep_prob = 1.0 - self.p
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+        random_tensor.floor_()
+        return x / keep_prob * random_tensor
+
 class Block(nn.Module):
     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=attn_drop, proj_drop=drop)
-        self.drop_path = nn.Identity() if drop_path == 0. else _StochasticDepth(drop_path)
+        self.drop_path = StochasticDepth(drop_path) if drop_path > 0. else nn.Identity()
         self.norm2 = nn.LayerNorm(dim)
-        self.mlp = MLP(dim, int(dim*mlp_ratio), drop=drop)
+        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__()
-        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
+        img_size = cfg["image_size"]
+
+        self.patch_embed = ConvPatchEmbed(
+            img_size=img_size,
+            in_chans=cfg["in_channels"],
+            embed_dim=cfg["emb_dim"]
+        )
+        n_patches = self.patch_embed.num_patches
 
-        self.cls_token = nn.Parameter(torch.zeros(1,1,cfg["emb_dim"]))
+        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_drop = nn.Dropout(p=cfg["drop"])
 
-        # transformer blocks
-        dpr = [x.item() for x in torch.linspace(0, cfg.get("drop_path", 0.2), cfg["depth"])]  # stochastic depth decay
+        # stochastic depth decay rule
+        dpr = torch.linspace(0, cfg["drop_path"], cfg["depth"]).tolist()
         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])
+            Block(
+                dim=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"])
@@ -174,9 +192,9 @@ class ViT(nn.Module):
 
     def forward(self, x):
         B = x.shape[0]
-        x = self.patch_embed(x)             # (B, N, E)
+        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 = torch.cat((cls_tokens, x), dim=1)     # (B, 1+N, E)
         x = x + self.pos_embed
         x = self.pos_drop(x)
 
@@ -190,7 +208,7 @@ class ViT(nn.Module):
 
 
 # Load model weights
-checkpoint = torch.load("best_CTD_ViT_CIFAR100_checkpoint.pth", map_location=device)
+checkpoint = torch.load("Revised_best_ViT_CIFAR100_baseline_checkpoint.pth", map_location=device)
 
 model = ViT(cfg).to(device)