Update efficientnet_transformer_glam.py

efficientnet_transformer_glam.py

@@ -15,6 +15,7 @@ def window_partition(x, window_size):
     windows = x.view(-1, window_size * window_size, C)
     return windows
 
+
 def window_reverse(windows, window_size, H, W):
     """Reverses the window partition operation."""
     B = int(windows.shape[0] / (H * W / window_size / window_size))
@@ -163,6 +164,9 @@ class EfficientNetb0_TransformerGLAM(nn.Module):
         # Self-adaptive gating
         self.gate_fc = nn.Linear(embed_dim, 1)
 
+        # Final feature output
+        self.final_feature_layer = nn.Identity()
+
         # Final classification
         self.dropout = nn.Dropout(dropout)
         self.fc = nn.Linear(embed_dim, num_classes)
@@ -183,7 +187,7 @@ class EfficientNetb0_TransformerGLAM(nn.Module):
         T = self.swin_attn(x_norm)
 
         T_perm = T.permute(0, 2, 3, 1).contiguous()
-        T_norm = self.post_attn_norm(T_perm)          # LN
+        T_norm = self.post_attn_norm(T_perm)           # LN
         T_out = T_norm.permute(0, 3, 1, 2).contiguous()
 
         # GLAM path
@@ -196,6 +200,8 @@ class EfficientNetb0_TransformerGLAM(nn.Module):
 
         # Final Fusion
         F_out = g * T_out + (1 - g) * G_out
+        F_out = self.final_feature_layer(F_out)  # ✅ Final feature map for Grad-CAM
         pooled = F.adaptive_avg_pool2d(F_out, (1, 1)).view(B, -1)
 
         return self.fc(pooled)
+