Upload 3 files

glam_efficientnet_model.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, PretrainedConfig, EfficientNetModel
+from typing import Optional, Union
+
+# --------------------------------------------------
+# Import your GLAM, SwinWindowAttention blocks here
+# --------------------------------------------------
+# from .glam_module import GLAM
+# from .swin_module import SwinWindowAttention
+
+from glam_module import GLAM
+from swin_module import SwinWindowAttention
+
+class GLAMEfficientNetConfig(PretrainedConfig):
+    """Hugging Face-style configuration for GLAM EfficientNet."""
+    model_type = "glam_efficientnet"
+
+    def __init__(self,
+                 num_classes: int = 3,
+                 embed_dim: int = 512,
+                 num_heads: int = 8,
+                 window_size: int = 7,
+                 reduction_ratio: int = 8,
+                 dropout: float = 0.5,
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.num_classes = num_classes
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.reduction_ratio = reduction_ratio
+        self.dropout = dropout
+
+
+class GLAMEfficientNetForClassification(PreTrainedModel):
+    """Hugging Face-style Model for EfficientNet + GLAM + Swin Architecture."""
+    config_class = GLAMEfficientNetConfig
+
+    def __init__(self, config: GLAMEfficientNetConfig, glam_module_cls, swin_module_cls):
+        super().__init__(config)
+
+        # ✅ 1) Hugging Face EfficientNet Backbone
+        self.features = EfficientNetModel.from_pretrained("google/efficientnet-b0")
+
+        # ✅ 1x1 conv for channel adjustment
+        self.conv1x1 = nn.Conv2d(1280, config.embed_dim, kernel_size=1)
+
+        # ✅ 2) Swin Attention Block
+        self.swin_attn = swin_module_cls(
+            embed_dim=config.embed_dim,
+            window_size=config.window_size,
+            num_heads=config.num_heads,
+            dropout=config.dropout
+        )
+        self.pre_attn_norm = nn.LayerNorm(config.embed_dim)
+        self.post_attn_norm = nn.LayerNorm(config.embed_dim)
+
+        # ✅ 3) GLAM Block
+        self.glam = glam_module_cls(in_channels=config.embed_dim, reduction_ratio=config.reduction_ratio)
+
+        # ✅ 4) Self-Adaptive Gating
+        self.gate_fc = nn.Linear(config.embed_dim, 1)
+
+        # ✅ Final classification
+        self.dropout = nn.Dropout(config.dropout)
+        self.classifier = nn.Linear(config.embed_dim, config.num_classes)
+
+    def forward(self, pixel_values, labels=None, **kwargs):
+        """Perform forward pass."""
+        # ✅ 1) EfficientNet Backbone
+        backbone_output = self.features(pixel_values)           # Returns BaseModelOutput
+        feats = backbone_output.last_hidden_state               # [B, C, H', W']
+        feats = self.conv1x1(feats)                             # Adjust channel dims
+        B, C, H, W = feats.shape
+
+        # ✅ 2) Transformer Branch
+        x_perm = feats.permute(0, 2, 3, 1).contiguous()         # [B, H', W', C]
+        x_norm = self.pre_attn_norm(x_perm).permute(0, 3, 1, 2).contiguous()
+        x_norm = self.dropout(x_norm)
+
+        T_out = self.swin_attn(x_norm)                           # [B, C, H', W']
+
+        T_out = self.post_attn_norm(T_out.permute(0, 2, 3, 1).contiguous())
+        T_out = T_out.permute(0, 3, 1, 2).contiguous()
+
+        # ✅ 3) GLAM Branch
+        G_out = self.glam(feats)
+
+        # ✅ 4) Self-Adaptive Gating
+        gap_feats = F.adaptive_avg_pool2d(feats, (1, 1)).view(B, C)
+        g = torch.sigmoid(self.gate_fc(gap_feats)).view(B, 1, 1, 1)
+
+        F_out = g * T_out + (1 - g) * G_out
+
+        # ✅ Final Pooling & Classifier
+        pooled = F.adaptive_avg_pool2d(F_out, (1, 1)).view(B, -1)
+        logits = self.classifier(self.dropout(pooled))
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits, labels)
+
+        return {"loss": loss, "logits": logits}
+

glam_module.py

@@ -0,0 +1,71 @@
+class GLAM(nn.Module):
+    """
+    Global-Local Attention Module (GLAM) that produces a refined feature map.
+    """
+    def __init__(self, in_channels, reduction_ratio=8):
+        super(GLAM, self).__init__()
+        
+        # --- Local Channel Attention ---
+        self.local_channel_conv = nn.Conv2d(in_channels, in_channels // reduction_ratio, kernel_size=1)
+        self.local_channel_act = nn.Sigmoid()
+        self.local_channel_expand = nn.Conv2d(in_channels // reduction_ratio, in_channels, kernel_size=1)
+        
+        # --- Local Spatial Attention ---
+        # 3-dilated, 5-dilated conv merges
+        self.local_spatial_conv3 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=3, dilation=3)
+        self.local_spatial_conv5 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=5, dilation=5)
+        self.local_spatial_merge = nn.Conv2d(in_channels * 3, in_channels, kernel_size=1)
+        self.local_spatial_act = nn.Sigmoid()
+        
+        # --- Global Channel Attention ---
+        self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.global_channel_fc1 = nn.Linear(in_channels, in_channels // reduction_ratio)
+        self.global_channel_fc2 = nn.Linear(in_channels // reduction_ratio, in_channels)
+        self.global_channel_act = nn.Sigmoid()
+        
+        # --- Global Spatial Attention ---
+        self.global_spatial_conv = nn.Conv2d(in_channels, 1, kernel_size=1)
+        self.global_spatial_softmax = nn.Softmax(dim=-1)
+        
+        
+        # --- Weighted paramerers initialization ---
+        self.local_attention_weight = nn.Parameter(torch.tensor(1.0))  
+        self.global_attention_weight = nn.Parameter(torch.tensor(1.0))
+        
+
+    def forward(self, x):
+        # Local Channel Attention
+        lca = self.local_channel_conv(x)       
+        lca = self.local_channel_act(lca)      
+        lca = self.local_channel_expand(lca)   
+        lca_out = lca * x                      
+
+        # Local Spatial Attention
+        lsa3 = self.local_spatial_conv3(x)
+        lsa5 = self.local_spatial_conv5(x)
+        lsa_cat = torch.cat([x, lsa3, lsa5], dim=1)
+        lsa = self.local_spatial_merge(lsa_cat)
+        lsa = self.local_spatial_act(lsa)
+        lsa_out = lsa * lca_out
+        lsa_out = lsa_out + lca_out
+
+        # Global Channel Attention
+        B, C, H, W = x.size()
+        gca = self.global_avg_pool(x).view(B, C)  
+        gca = F.relu(self.global_channel_fc1(gca), inplace=True)
+        gca = self.global_channel_fc2(gca)
+        gca = self.global_channel_act(gca)
+        gca = gca.view(B, C, 1, 1)
+        gca_out = gca * x
+
+        # Global Spatial Attention
+        gsa = self.global_spatial_conv(x)  # [B, 1, H, W]
+        gsa = gsa.view(B, -1)              # [B, H*W]
+        gsa = self.global_spatial_softmax(gsa)
+        gsa = gsa.view(B, 1, H, W)
+        gsa_out = gsa * gca_out
+        gsa_out = gsa_out + gca_out
+
+        # Fuse
+        out =  lsa_out*self.local_attention_weight +  gsa_out*self.global_attention_weight + x
+        return out

swin_module.py

@@ -0,0 +1,72 @@
+# -------------------------------
+# 2. SWIN-STYLE TRANSFORMER UTILS
+# -------------------------------
+def window_partition(x, window_size):
+    """
+    x: (B, H, W, C)
+    Returns windows of shape: (num_windows*B, window_size*window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    # permute to gather patches
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
+    # merge dimension
+    windows = x.view(-1, window_size * window_size, C)
+    return windows
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Reverse of window_partition.
+    windows: (num_windows*B, window_size*window_size, C)
+    Returns: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous()
+    x = x.view(B, H, W, -1)
+    return x
+
+class SwinWindowAttention(nn.Module):
+    """
+    A simplified Swin-like window attention block:
+    1) Partition input into windows
+    2) Perform multi-head self-attn
+    3) Merge back
+    """
+    def __init__(self, embed_dim, window_size, num_heads, dropout=0.0):
+        super(SwinWindowAttention, self).__init__()
+        self.embed_dim = embed_dim
+        self.window_size = window_size
+        self.num_heads = num_heads
+        
+        self.mha = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, batch_first=True)
+        self.dropout = nn.Dropout(dropout)
+    
+    def forward(self, x):
+        # x: (B, C, H, W) --> rearrange to (B, H, W, C)
+        B, C, H, W = x.shape
+        x = x.permute(0, 2, 3, 1).contiguous()
+        
+        # pad if needed so H, W are multiples of window_size
+        pad_h = (self.window_size - H % self.window_size) % self.window_size
+        pad_w = (self.window_size - W % self.window_size) % self.window_size
+        if pad_h or pad_w:
+            x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
+        
+        Hp, Wp = x.shape[1], x.shape[2]
+        # Partition into windows
+        windows = window_partition(x, self.window_size)  # shape: (num_windows*B, window_size*window_size, C)
+        # Multi-head self-attn
+        attn_windows, _ = self.mha(windows, windows, windows)
+        attn_windows = self.dropout(attn_windows)
+        
+        # Reverse window partition
+        x = window_reverse(attn_windows, self.window_size, Hp, Wp)
+        
+        # Remove padding if added
+        if pad_h or pad_w:
+            x = x[:, :H, :W, :].contiguous()
+        
+        # back to (B, C, H, W)
+        x = x.permute(0, 3, 1, 2).contiguous()
+        return x