Update efficientnet_transformer_glam.py

efficientnet_transformer_glam.py

@@ -4,9 +4,9 @@ import torch.nn.functional as F
 import torchvision.models as models
 
 
-# -------------------------------
+# ==================================================
 # 1. SWIN WINDOW UTILS
-# -------------------------------
+# ==================================================
 def window_partition(x, window_size):
     """Partitions input tensor into windows of shape (B * num_windows, window_size*window_size, C)."""
     B, H, W, C = x.shape
@@ -25,9 +25,9 @@ def window_reverse(windows, window_size, H, W):
     return x
 
 
-# -------------------------------
+# ==================================================
 # 2. SWIN WINDOW ATTENTION
-# -------------------------------
+# ==================================================
 class SwinWindowAttention(nn.Module):
     """Swin-style window attention block."""
     def __init__(self, embed_dim, window_size, num_heads, dropout=0.0):
@@ -62,9 +62,9 @@ class SwinWindowAttention(nn.Module):
         return x.permute(0, 3, 1, 2).contiguous()
 
 
-# -------------------------------
+# ==================================================
 # 3. GLAM
-# -------------------------------
+# ==================================================
 class GLAM(nn.Module):
     """Global-Local Attention Module (GLAM)."""
     def __init__(self, in_channels, reduction_ratio=8):
@@ -133,9 +133,22 @@ class GLAM(nn.Module):
         return out
 
 
-# -------------------------------
-# 4. EFFICIENTNETB0_TRANSFORMERGLAM
-# -------------------------------
+# ==================================================
+# 4. FUSION BLOCK
+# ==================================================
+class FusionBlock(nn.Module):
+    """Combines Transformer and GLAM outputs using gating."""
+    def __init__(self):
+        super(FusionBlock, self).__init__()
+
+    def forward(self, g, T_out, G_out):
+        """Perform final gating fusion."""
+        return g * T_out + (1 - g) * G_out
+
+
+# ==================================================
+# 5. EFFICIENTNETB0_TRANSFORMERGLAM
+# ==================================================
 class EfficientNetb0_TransformerGLAM(nn.Module):
     """EfficientNet-B0 + Swin-style Transformer + GLAM + Self-Adaptive Gating."""
     def __init__(self,
@@ -161,11 +174,11 @@ class EfficientNetb0_TransformerGLAM(nn.Module):
         # GLAM path
         self.glam = GLAM(in_channels=embed_dim, reduction_ratio=reduction_ratio)
 
-        # Self-adaptive gating
+        # Gating
         self.gate_fc = nn.Linear(embed_dim, 1)
 
-        # Final feature output
-        self.final_feature_layer = nn.Identity()
+        # Final Fusion
+        self.fusion_block = FusionBlock()
 
         # Final classification
         self.dropout = nn.Dropout(dropout)
@@ -199,9 +212,12 @@ class EfficientNetb0_TransformerGLAM(nn.Module):
         g = g.view(B, 1, 1, 1)
 
         # Final Fusion
-        F_out = g * T_out + (1 - g) * G_out
-        # ✅ Save the spatial feature map
+        F_out = self.fusion_block(g, T_out, G_out)
+
+        # Save final feature map for Grad-CAM
         self.last_feature = F_out
+
         pooled = F.adaptive_avg_pool2d(F_out, (1, 1)).view(B, -1)
+
         return self.fc(pooled)