Update glam_efficientnet_model.py

glam_efficientnet_model.py

@@ -1,106 +1,95 @@
-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}
-
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import models
+from typing import Optional, Union
+
+from glam_module import GLAM
+from swin_module import SwinWindowAttention
+
+
+class GLAMEfficientNetConfig:
+    """Hugging Face-style configuration for 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(nn.Module):
+    """EfficientNet (torchvision) + GLAM + Swin Architecture for Classification."""
+    def __init__(self, config: GLAMEfficientNetConfig, glam_module_cls, swin_module_cls):
+        super().__init__()
+
+        # ✅ 1) Torchvision EfficientNet Backbone
+        efficientnet = models.efficientnet_b0(pretrained=False)  # No Hugging Face!
+        self.features = efficientnet.features
+
+        # ✅ 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
+        feats = self.features(pixel_values)              # [B, 1280, H', W']
+        feats = self.conv1x1(feats)                      # [B, embed_dim, H', W']
+        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}