updated model.py

model.py

@@ -1,7 +1,6 @@
 import torch
-import torchvision
 import torch.nn as nn
-import torch.optim as optim
+import torchvision
 import math
 
 class ImprovedEfficientBackbone(nn.Module):
@@ -11,8 +10,7 @@ class ImprovedEfficientBackbone(nn.Module):
         self.features = self.efficientnet.features
 
     def forward(self, x):
-        return self.features(x) 
-
+        return self.features(x)  
 
 class ImprovedPatchEmbedding(nn.Module):
     def __init__(self, in_channels=1280, embed_dim=384):
@@ -25,7 +23,7 @@ class ImprovedPatchEmbedding(nn.Module):
         Output: [B, 49, 384]
         """
         B, C, H, W = x.shape
-        x = x.flatten(2).transpose(1, 2) 
+        x = x.flatten(2).transpose(1, 2)  
         x = self.proj(x)  
         return x
 
@@ -41,23 +39,24 @@ class ImprovedViTBlock(nn.Module):
             nn.GELU(),
             nn.Linear(embed_dim * mlp_ratio, embed_dim)
         )
-        self.dropout = nn.Dropout(0.1)
+        self.dropout = nn.Dropout(0.2) 
 
     def forward(self, x):
         x = x + self.dropout(self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0])
         x = x + self.dropout(self.mlp(self.norm2(x)))
         return x
 
-
-
 class ImprovedEfficientViT(nn.Module):
-    def __init__(self, embed_dim=384, depth=8, num_heads=4):
+    def __init__(self, embed_dim=384, depth=6, num_heads=4):
         super().__init__()
         self.backbone = ImprovedEfficientBackbone()
         self.patch_embed = ImprovedPatchEmbedding(embed_dim=embed_dim)
 
         self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))
-        self.register_buffer("pos_embed", self._get_sinusoidal_encoding(50, embed_dim))  # Use sin/cos
+        self.register_buffer("pos_embed", self._get_sinusoidal_encoding(50, embed_dim))
+
+        self.patch_dropout = nn.Dropout(0.2)
+        self.pos_dropout = nn.Dropout(0.2)
 
         self.blocks = nn.ModuleList([ImprovedViTBlock(embed_dim, num_heads) for _ in range(depth)])
 
@@ -65,6 +64,7 @@ class ImprovedEfficientViT(nn.Module):
               nn.LayerNorm(embed_dim),
               nn.Linear(embed_dim, 128),
               nn.GELU(),
+              nn.Dropout(0.3),
               nn.Linear(128, 1)
           )
 
@@ -82,16 +82,18 @@ class ImprovedEfficientViT(nn.Module):
         return pe.unsqueeze(0)  
 
     def forward(self, x):
-        features = self.backbone(x)                
-        tokens = self.patch_embed(features)         
+        features = self.backbone(x)                  
+        tokens = self.patch_embed(features)          
+        tokens = self.patch_dropout(tokens)          
 
         B = tokens.shape[0]
-        cls_tokens = self.cls_token.expand(B, -1, -1) 
+        cls_tokens = self.cls_token.expand(B, -1, -1)  
         x = torch.cat((cls_tokens, tokens), dim=1)    
         x = x + self.pos_embed[:, :x.size(1), :]      
+        x = self.pos_dropout(x)                       
 
         for block in self.blocks:
             x = block(x)
 
-        cls_final = x[:, 0]  
-        return self.head(cls_final) 
+        cls_final = x[:, 0] 
+        return self.head(cls_final)