Spaces:
Sleeping
Sleeping
File size: 5,535 Bytes
8c58a75 | 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 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """
Model Loader — supports multiple pretrained torchvision architectures
with helpers for extracting named convolutional layers for Grad-CAM targeting.
"""
import torch
import torch.nn as nn
import torchvision.models as models
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple
@dataclass
class ModelConfig:
"""Configuration for a pretrained vision model."""
name: str
constructor: callable
weights_class: object
default_target_layer: str # Attribute path string for the default Grad-CAM layer
cam_layers: List[str] # All available layers for Grad-CAM
description: str
# Registry of supported models
MODEL_REGISTRY: Dict[str, ModelConfig] = {
"ResNet-50": ModelConfig(
name="ResNet-50",
constructor=models.resnet50,
weights_class=models.ResNet50_Weights.IMAGENET1K_V2,
default_target_layer="layer4.2.conv3",
cam_layers=["layer1", "layer2", "layer3", "layer4", "layer4.2.conv3"],
description="50-layer residual network. Great balance of speed and accuracy.",
),
"ResNet-101": ModelConfig(
name="ResNet-101",
constructor=models.resnet101,
weights_class=models.ResNet101_Weights.IMAGENET1K_V2,
default_target_layer="layer4.2.conv3",
cam_layers=["layer1", "layer2", "layer3", "layer4", "layer4.2.conv3"],
description="101-layer residual network. Higher accuracy, slower inference.",
),
"VGG-16": ModelConfig(
name="VGG-16",
constructor=models.vgg16,
weights_class=models.VGG16_Weights.IMAGENET1K_V1,
default_target_layer="features.28",
cam_layers=["features.14", "features.21", "features.28"],
description="Classic VGG architecture. Very clear gradient signals.",
),
"EfficientNet-B0": ModelConfig(
name="EfficientNet-B0",
constructor=models.efficientnet_b0,
weights_class=models.EfficientNet_B0_Weights.IMAGENET1K_V1,
default_target_layer="features.8.0",
cam_layers=["features.3", "features.5", "features.7", "features.8.0"],
description="Efficient compound-scaled architecture. Small and fast.",
),
"DenseNet-121": ModelConfig(
name="DenseNet-121",
constructor=models.densenet121,
weights_class=models.DenseNet121_Weights.IMAGENET1K_V1,
default_target_layer="features.denseblock4.denselayer16.conv2",
cam_layers=["features.denseblock2", "features.denseblock3", "features.denseblock4"],
description="Dense connections between all layers. Strong feature reuse.",
),
"MobileNet-V3": ModelConfig(
name="MobileNet-V3",
constructor=models.mobilenet_v3_large,
weights_class=models.MobileNet_V3_Large_Weights.IMAGENET1K_V2,
default_target_layer="features.16.block.2.0",
cam_layers=["features.7", "features.12", "features.16.block.2.0"],
description="Lightweight mobile-optimized network.",
),
}
def load_model(model_name: str, device: Optional[str] = None) -> Tuple[nn.Module, ModelConfig]:
"""
Load a pretrained model by name from the registry.
Args:
model_name: Key in MODEL_REGISTRY
device: 'cpu', 'cuda', or None (auto-detect)
Returns:
(model, config) tuple
"""
if model_name not in MODEL_REGISTRY:
raise ValueError(f"Unknown model '{model_name}'. Available: {list(MODEL_REGISTRY.keys())}")
config = MODEL_REGISTRY[model_name]
if device is None:
device = "cuda" if torch.cuda.is_available() else "cpu"
model = config.constructor(weights=config.weights_class)
model = model.to(device)
model.eval()
# Disable inplace ops (VGG/EfficientNet/MobileNet use inplace ReLU
# which causes autograd view errors with backward hooks)
for module in model.modules():
if hasattr(module, 'inplace'):
module.inplace = False
return model, config
def get_layer_by_name(model: nn.Module, layer_name: str) -> nn.Module:
"""
Navigate the model hierarchy using a dot-separated attribute path.
E.g., 'layer4.2.conv3' → model.layer4[2].conv3
Handles both attribute access and integer indexing for Sequential/ModuleList.
"""
parts = layer_name.split(".")
current = model
for part in parts:
if part.isdigit():
current = current[int(part)]
else:
current = getattr(current, part)
return current
def list_conv_layers(model: nn.Module, prefix: str = "") -> List[Tuple[str, nn.Module]]:
"""
Recursively list all Conv2d layers in the model with their full attribute paths.
Useful for letting users explore all possible Grad-CAM target layers.
"""
conv_layers = []
for name, module in model.named_modules():
if isinstance(module, nn.Conv2d):
conv_layers.append((name, module))
return conv_layers
def get_model_summary(model: nn.Module) -> Dict:
"""Return a summary dict with parameter counts and layer info."""
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
conv_count = sum(1 for m in model.modules() if isinstance(m, nn.Conv2d))
bn_count = sum(1 for m in model.modules() if isinstance(m, (nn.BatchNorm2d,)))
return {
"total_params": total_params,
"trainable_params": trainable_params,
"conv_layers": conv_count,
"batch_norm_layers": bn_count,
}
|