Spaces:
Sleeping
Sleeping
File size: 2,662 Bytes
af59988 |
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 |
"""
EfficientNet-B0 model for Pneumonia classification.
"""
import torch
import torch.nn as nn
from torchvision import models
from typing import Tuple
from .config import DROPOUT_RATE, NUM_CLASSES
class PneumoniaClassifier(nn.Module):
"""EfficientNet-B0 based classifier for chest X-ray pneumonia detection."""
def __init__(
self,
pretrained: bool = True,
dropout_rate: float = DROPOUT_RATE,
freeze_backbone: bool = True
):
super().__init__()
# Load pretrained EfficientNet-B0
weights = models.EfficientNet_B0_Weights.IMAGENET1K_V1 if pretrained else None
self.backbone = models.efficientnet_b0(weights=weights)
# Get the number of features from the classifier
in_features = self.backbone.classifier[1].in_features # 1280
# Replace classifier head
self.backbone.classifier = nn.Sequential(
nn.Dropout(p=dropout_rate, inplace=True),
nn.Linear(in_features, NUM_CLASSES)
)
# Freeze backbone if specified
if freeze_backbone:
self.freeze_backbone()
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.backbone(x)
def freeze_backbone(self):
"""Freeze all layers except the classifier."""
for param in self.backbone.features.parameters():
param.requires_grad = False
def unfreeze_backbone(self):
"""Unfreeze all layers for fine-tuning."""
for param in self.backbone.features.parameters():
param.requires_grad = True
def get_param_counts(self) -> Tuple[int, int]:
"""Return (trainable_params, total_params)."""
trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
total = sum(p.numel() for p in self.parameters())
return trainable, total
def create_model(
pretrained: bool = True,
dropout_rate: float = DROPOUT_RATE,
freeze_backbone: bool = True,
device: str = None
) -> PneumoniaClassifier:
"""Factory function to create the model."""
if device is None:
device = "mps" if torch.backends.mps.is_available() else \
"cuda" if torch.cuda.is_available() else "cpu"
model = PneumoniaClassifier(
pretrained=pretrained,
dropout_rate=dropout_rate,
freeze_backbone=freeze_backbone
)
return model.to(device)
def get_device() -> torch.device:
"""Get the best available device."""
if torch.backends.mps.is_available():
return torch.device("mps")
elif torch.cuda.is_available():
return torch.device("cuda")
return torch.device("cpu")
|