Spaces:
Runtime error
Runtime error
File size: 1,585 Bytes
78856c0 | 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 | """Backbone feature extractors supporting EfficientNet variants via timm."""
from typing import Tuple
import timm
import torch
import torch.nn as nn
class EfficientNetBackbone(nn.Module):
"""EfficientNet feature extractor wrapping timm pretrained models."""
def __init__(
self,
model_name: str = "efficientnet_b0",
pretrained: bool = True,
drop_rate: float = 0.2,
drop_path_rate: float = 0.1,
) -> None:
super().__init__()
self.model_name = model_name
self.pretrained = pretrained
# Create model with head removed (num_classes=0 returns pooled feature vector)
self.encoder = timm.create_model(
model_name,
pretrained=pretrained,
num_classes=0,
drop_rate=drop_rate,
drop_path_rate=drop_path_rate,
)
self.num_features = self.encoder.num_features
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Extract global pooled feature representations.
Args:
x: Input tensor of shape (N, C, H, W).
Returns:
Feature embedding of shape (N, num_features).
"""
return self.encoder(x)
def forward_features(self, x: torch.Tensor) -> torch.Tensor:
"""Extract unpooled spatial feature maps before global pooling (for Grad-CAM).
Args:
x: Input tensor of shape (N, C, H, W).
Returns:
Spatial feature maps of shape (N, C_feat, H_feat, W_feat).
"""
return self.encoder.forward_features(x)
|