File size: 4,093 Bytes
d10a3b5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
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}
|