File size: 3,729 Bytes
63f5b1f | 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 | from __future__ import annotations
import torch
import torch.nn as nn
import timm
class HybridDN121EffB0(nn.Module):
def __init__(self, num_classes=2, head_dim=256, dropout=0.2):
super().__init__()
self.bb1 = timm.create_model("densenet121", pretrained=False, num_classes=0, global_pool="avg")
self.bb2 = timm.create_model("efficientnet_b0", pretrained=False, num_classes=0, global_pool="avg")
self.head = nn.Sequential(nn.Linear(self.bb1.num_features + self.bb2.num_features, head_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(head_dim, num_classes))
def forward(self, x):
return self.head(torch.cat([self.bb1(x), self.bb2(x)], dim=1))
class HybridSwinTEffB0(nn.Module):
def __init__(self, num_classes=2, head_dim=256, dropout=0.2):
super().__init__()
self.bb1 = timm.create_model("swin_tiny_patch4_window7_224", pretrained=False, num_classes=0, global_pool="avg")
self.bb2 = timm.create_model("efficientnet_b0", pretrained=False, num_classes=0, global_pool="avg")
self.head = nn.Sequential(nn.Linear(self.bb1.num_features + self.bb2.num_features, head_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(head_dim, num_classes))
def forward(self, x):
return self.head(torch.cat([self.bb1(x), self.bb2(x)], dim=1))
class SEBlock(nn.Module):
def __init__(self, channels, reduction=16):
super().__init__()
hidden = max(8, channels // reduction)
self.pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Sequential(nn.Linear(channels, hidden, bias=False), nn.GELU(), nn.Linear(hidden, channels, bias=False), nn.Sigmoid())
def forward(self, x):
b, c, _, _ = x.shape
return x * self.fc(self.pool(x).view(b, c)).view(b, c, 1, 1)
class Conv1x1BNAct(nn.Module):
def __init__(self, in_ch, out_ch):
super().__init__()
self.net = nn.Sequential(nn.Conv2d(in_ch, out_ch, 1, bias=False), nn.BatchNorm2d(out_ch), nn.GELU())
def forward(self, x):
return self.net(x)
class CustomMSAFEffB0(nn.Module):
def __init__(self, num_classes=2, embed_dim=256, head_dim=256, dropout=0.35):
super().__init__()
self.backbone = timm.create_model("efficientnet_b0", pretrained=False, features_only=True, out_indices=(2, 3, 4))
channels = self.backbone.feature_info.channels()
self.proj = nn.ModuleList([Conv1x1BNAct(c, embed_dim) for c in channels])
self.se = nn.ModuleList([SEBlock(embed_dim) for _ in channels])
self.pool = nn.AdaptiveAvgPool2d(1)
self.scale_attn = nn.Sequential(nn.Linear(embed_dim, max(8, embed_dim // 4)), nn.GELU(), nn.Dropout(dropout * 0.25), nn.Linear(max(8, embed_dim // 4), 1))
self.head = nn.Sequential(nn.Linear(embed_dim * (len(channels) + 1), head_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(head_dim, num_classes))
def forward(self, x):
vecs = [self.pool(se(proj(f))).flatten(1) for f, proj, se in zip(self.backbone(x), self.proj, self.se)]
ms = torch.stack(vecs, dim=1)
weights = torch.softmax(self.scale_attn(ms).squeeze(-1), dim=1)
return self.head(torch.cat([(ms * weights.unsqueeze(-1)).sum(dim=1), ms.flatten(1)], dim=1))
def build_model(architecture, num_classes=2):
if architecture == "hybrid_dn121_effb0":
return HybridDN121EffB0(num_classes=num_classes)
if architecture == "hybrid_swint_effb0":
return HybridSwinTEffB0(num_classes=num_classes)
if architecture == "custom_msaf_effb0":
return CustomMSAFEffB0(num_classes=num_classes)
return timm.create_model(architecture, pretrained=False, num_classes=num_classes)
|