Upload LW_VIT.py

LW_VIT.py

@@ -0,0 +1,181 @@
+"""LW-Vit"""
+
+import torch 
+import torch.nn as nn
+import torch.nn.functional as F
+
+class HSwish(nn.Module):
+    """Hard Swish activation function"""
+    def forward(self, x):
+        out = x * F.relu6(x + 3, inplace=True) / 6
+        return out
+
+class MV2_Block(nn.Module):
+    """MobileNetV2 Inverted Residual Block with h-swish activation"""
+    def __init__(self, in_channels, expand_channels, out_channels, stride):
+        super().__init__()
+        self.stride = stride
+        self.use_res_connect = self.stride == 1 and in_channels == out_channels
+        
+        self.expand = nn.Sequential(
+            nn.Conv2d(in_channels, expand_channels, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.BatchNorm2d(expand_channels),
+            HSwish()
+        )
+
+        self.depthwise = nn.Sequential(
+            nn.Conv2d(expand_channels, expand_channels, kernel_size=3, stride=stride, 
+                     padding=1, groups=expand_channels, bias=False),
+            nn.BatchNorm2d(expand_channels),
+            HSwish()
+        )
+
+        self.project = nn.Sequential(
+            nn.Conv2d(expand_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.BatchNorm2d(out_channels)
+        )
+
+    def forward(self, x):
+        if self.use_res_connect:
+            return x + self.project(self.depthwise(self.expand(x)))
+        else:
+            return self.project(self.depthwise(self.expand(x)))
+        
+class LW_ViT_Transformer_Block(nn.Module):
+    """Lightweight Vision Transformer Block"""
+    def __init__(self, in_channels, out_channels, patch_size=2, heads=4, dim=128):
+        super().__init__()
+        self.patch_size = patch_size
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.dim = dim
+        
+        # Downsampling and channel adjustment
+        self.downsample = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=1)
+        self.channel_adjust = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1)
+        
+        # Calculate patch dimension
+        patch_dim = out_channels * patch_size * patch_size
+        # Add projection layer if needed
+        self.projection = nn.Identity() if patch_dim == dim else nn.Linear(patch_dim, dim)
+        
+        # Transformer components
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = nn.MultiheadAttention(embed_dim=dim, num_heads=heads, batch_first=True)
+        self.norm2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim)
+        )
+        
+        # Final processing
+        self.final_conv = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
+        self.norm_out = nn.BatchNorm2d(out_channels)
+        self.act = HSwish()
+
+    def forward(self, x):
+        # Initial downsampling and channel adjustment
+        x = self.downsample(x)
+        x = self.channel_adjust(x)
+        
+        B, C, H, W = x.shape
+        
+        # Convert feature map to patches and embed
+        # Reshape to [B, C, H/P, P, W/P, P]
+        x_reshaped = x.reshape(B, C, H // self.patch_size, self.patch_size, 
+                              W // self.patch_size, self.patch_size)
+        # Permute to [B, H/P, W/P, C, P, P]
+        x_permuted = x_reshaped.permute(0, 2, 4, 1, 3, 5)
+        # Reshape to [B, H/P * W/P, C * P * P]
+        patches = x_permuted.reshape(B, (H // self.patch_size) * (W // self.patch_size), -1)
+        
+        # Apply projection if needed
+        patches = self.projection(patches)
+        
+        # Apply transformer operations
+        normed_patches = self.norm1(patches)
+        attn_out, _ = self.attn(normed_patches, normed_patches, normed_patches)
+        patches = patches + attn_out
+        
+        normed_patches = self.norm2(patches)
+        ffn_out = self.ffn(normed_patches)
+        patches = patches + ffn_out
+        
+        # Reshape back to feature map
+        # Reshape to [B, H/P, W/P, C, P, P]
+        x_back = patches.reshape(B, H // self.patch_size, W // self.patch_size, 
+                                C, self.patch_size, self.patch_size)
+        # Permute to [B, C, H/P, P, W/P, P]
+        x_back = x_back.permute(0, 3, 1, 4, 2, 5)
+        # Reshape to [B, C, H, W]
+        x_out = x_back.reshape(B, C, H, W)
+        
+        # Final processing
+        out = self.final_conv(x_out)
+        out = self.norm_out(out)
+        out = self.act(out)
+        
+        return out
+
+class LW_ViT(nn.Module):
+    """Lightweight Vision Transformer for Chinese Character Recognition"""
+    def __init__(self, base_channels=16, num_classes=3892):
+        super().__init__()
+        
+        # Stem layer
+        self.stem = nn.Sequential(
+            nn.Conv2d(3, base_channels, 3, stride=2, padding=1, bias=False),
+            nn.BatchNorm2d(base_channels),
+            HSwish()
+        )
+
+        # Feature extraction layers
+        self.features = nn.Sequential(
+            # MV2 blocks
+            MV2_Block(base_channels, 4*base_channels, base_channels, stride=1),
+            MV2_Block(base_channels, 4*base_channels, base_channels, stride=2),
+            MV2_Block(base_channels, 4*base_channels, base_channels, stride=1),
+            MV2_Block(base_channels, 4*base_channels, base_channels, stride=2),
+            
+            # Transformer block
+            LW_ViT_Transformer_Block(
+                in_channels=base_channels, 
+                out_channels=2*base_channels, 
+                patch_size=2, 
+                heads=4, 
+                dim=128
+            ),
+            
+            # Final convolutional layer
+            nn.Conv2d(2*base_channels, 4*base_channels, kernel_size=1, stride=1),
+            nn.BatchNorm2d(4*base_channels),
+            HSwish()
+        )
+        
+        # Global pooling and classifier
+        self.global_pool = nn.AdaptiveAvgPool2d(1)
+        self.dropout = nn.Dropout(0.1)  # Added dropout for regularization
+        
+        self.classifier = nn.Sequential(
+            nn.Linear(4*base_channels, 8*base_channels),
+            nn.BatchNorm1d(8*base_channels),
+            HSwish(),
+            nn.Dropout(0.1),  # Added dropout for regularization
+            nn.Linear(8*base_channels, num_classes)
+        )
+
+    def forward(self, x):
+        x = self.stem(x)
+        x = self.features(x)
+        x = self.global_pool(x)
+        x = torch.flatten(x, 1)
+        x = self.dropout(x)
+        x = self.classifier(x)
+        return x
+    
+
+num_cls = 3892
+num_epochs = 100
+user_device = "cuda" if torch.cuda.is_available() else "cpu"
+LW_model = LW_ViT(base_channels=16, num_classes=num_cls)