File size: 5,178 Bytes
565aecf | 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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | """
model.py
--------
EfficientNetB3-based leaf disease classifier.
- Pretrained on ImageNet for strong generalization
- Custom classification head
- Supports feature extraction + fine-tuning phases
"""
import torch
import torch.nn as nn
import timm
class LeafDiseaseModel(nn.Module):
"""
EfficientNetB3 with custom classification head for leaf disease detection.
Architecture:
EfficientNetB3 backbone (ImageNet pretrained)
→ Global Average Pooling
→ BatchNorm → Dense(512) → GELU → Dropout(0.4)
→ BatchNorm → Dense(256) → GELU → Dropout(0.3)
→ Dense(num_classes) → Softmax
"""
def __init__(self, num_classes: int, pretrained: bool = True, dropout: float = 0.4):
super().__init__()
self.num_classes = num_classes
# ── Backbone ──────────────────────────────────────────────────────────
self.backbone = timm.create_model(
"efficientnet_b3",
pretrained=pretrained,
num_classes=0, # remove default head
global_pool="avg",
)
backbone_out = self.backbone.num_features # 1536 for B3
# ── Custom Head ───────────────────────────────────────────────────────
self.head = nn.Sequential(
nn.BatchNorm1d(backbone_out),
nn.Linear(backbone_out, 512),
nn.GELU(),
nn.Dropout(dropout),
nn.BatchNorm1d(512),
nn.Linear(512, 256),
nn.GELU(),
nn.Dropout(dropout * 0.75),
nn.Linear(256, num_classes),
)
# Weight initialisation for the head
for m in self.head.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight)
if m.bias is not None:
nn.init.zeros_(m.bias)
# ── Phase control ─────────────────────────────────────────────────────────
def freeze_backbone(self):
"""Freeze backbone — only train the head (Phase 1)."""
for p in self.backbone.parameters():
p.requires_grad = False
print(" Backbone frozen — training head only.")
def unfreeze_backbone(self, unfreeze_layers: int = 30):
"""
Unfreeze the last N backbone layers for fine-tuning (Phase 2).
EfficientNetB3 has ~360 parameters groups; last 30 covers blocks 5-7.
"""
all_params = list(self.backbone.parameters())
# First, freeze everything
for p in all_params:
p.requires_grad = False
# Then unfreeze last N
for p in all_params[-unfreeze_layers:]:
p.requires_grad = True
trainable = sum(p.numel() for p in self.backbone.parameters() if p.requires_grad)
print(f" Unfrozen last {unfreeze_layers} backbone param groups "
f"({trainable:,} params now trainable).")
def unfreeze_all(self):
"""Fully unfreeze everything (Phase 3)."""
for p in self.parameters():
p.requires_grad = True
print(" All layers unfrozen.")
# ── Forward ───────────────────────────────────────────────────────────────
def forward(self, x: torch.Tensor) -> torch.Tensor:
features = self.backbone(x) # (B, 1536)
logits = self.head(features) # (B, num_classes)
return logits
def get_probabilities(self, x: torch.Tensor) -> torch.Tensor:
"""Return softmax probabilities."""
return torch.softmax(self.forward(x), dim=-1)
def predict(self, x: torch.Tensor):
"""Return (class_idx, confidence) tuple."""
probs = self.get_probabilities(x)
conf, idx = torch.max(probs, dim=-1)
return idx, conf
# ── Utilities ─────────────────────────────────────────────────────────────
def count_parameters(self):
total = sum(p.numel() for p in self.parameters())
trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
print(f" Total parameters: {total:>12,}")
print(f" Trainable parameters: {trainable:>12,}")
return total, trainable
def build_model(num_classes: int, pretrained: bool = True) -> LeafDiseaseModel:
"""Factory function — builds and returns the model."""
model = LeafDiseaseModel(num_classes=num_classes, pretrained=pretrained)
return model
if __name__ == "__main__":
m = build_model(39)
m.count_parameters()
x = torch.randn(4, 3, 300, 300)
out = m(x)
print(f" Output shape: {out.shape}") # (4, 39)
|