Spaces:
Sleeping
Sleeping
File size: 7,566 Bytes
fc33673 d7e388e 8645c1a d7e388e fc33673 7620674 fc33673 d7e388e 7620674 fc33673 d7e388e fc33673 d7e388e fc33673 d7e388e fc33673 d7e388e fc33673 d7e388e fc33673 7620674 fc33673 7620674 fc33673 d7e388e 8645c1a d7e388e 8645c1a d7e388e fc33673 d7e388e fc33673 d7e388e | 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 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | from __future__ import annotations
from typing import Dict
import torch
import timm
from torchvision.models import (
ConvNeXt_Tiny_Weights,
EfficientNet_B0_Weights,
MobileNet_V3_Small_Weights,
ResNet18_Weights,
Swin_T_Weights,
ViT_B_16_Weights,
convnext_tiny,
efficientnet_b0,
mobilenet_v3_small,
resnet18,
swin_t,
vit_b_16,
)
from src.models.lora import (
apply_lora_to_timm_attention,
enable_lora_suffix_no_grad,
freeze_all_parameters_except_lora_and_heads,
)
_TIMM_CLASSIFIER_MODELS = {
"vit_s16_cls": "vit_small_patch16_224",
"vit_l16_cls": "vit_large_patch16_224",
}
def _resolve_timm_classifier_name(model_cfg: Dict) -> str | None:
model_name = str(model_cfg.get("name", "resnet18_cls")).lower()
default_timm_name = _TIMM_CLASSIFIER_MODELS.get(model_name)
if default_timm_name is None:
return None
return str(model_cfg.get("timm_model_name", default_timm_name))
def _pretrained_weights(pretrained: bool, weights_enum):
return weights_enum.DEFAULT if pretrained else None
def _make_head(in_features: int, num_classes: int, dropout_p: float) -> torch.nn.Module:
if dropout_p > 0.0:
return torch.nn.Sequential(
torch.nn.Dropout(p=dropout_p),
torch.nn.Linear(in_features, num_classes),
)
return torch.nn.Linear(in_features, num_classes)
class MultiHeadResNet18Classifier(torch.nn.Module):
def __init__(
self,
pretrained: bool,
num_classes: int,
aux_dims: dict[str, int],
dropout_p: float,
) -> None:
super().__init__()
self.backbone = resnet18(weights=_pretrained_weights(pretrained, ResNet18_Weights))
in_features = self.backbone.fc.in_features
self.backbone.fc = torch.nn.Identity()
self.class_head = _make_head(in_features, num_classes, dropout_p)
self.brand_head = _make_head(in_features, aux_dims["brand"], dropout_p)
self.year_head = _make_head(in_features, aux_dims["year"], dropout_p)
def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
features = self.backbone(x)
return {
"class": self.class_head(features),
"brand": self.brand_head(features),
"year": self.year_head(features),
}
def get_classifier_data_config(model_cfg: Dict) -> dict:
timm_model_name = _resolve_timm_classifier_name(model_cfg)
if timm_model_name is None:
return {}
pretrained_cfg = timm.get_pretrained_cfg(timm_model_name)
if pretrained_cfg is None:
return {}
return {
"source": "timm_pretrained_cfg",
"input_size": list(pretrained_cfg.input_size),
"interpolation": pretrained_cfg.interpolation,
"crop_pct": float(pretrained_cfg.crop_pct),
"normalization": {
"mean": [float(v) for v in pretrained_cfg.mean],
"std": [float(v) for v in pretrained_cfg.std],
},
}
def build_classifier(model_cfg: Dict, num_classes: int, aux_dims: dict[str, int] | None = None):
model_name = str(model_cfg.get("name", "resnet18_cls")).lower()
pretrained = bool(model_cfg.get("pretrained", False))
dropout_p = float(model_cfg.get("dropout", 0.0))
use_aux_heads = bool(model_cfg.get("aux_heads", {}).get("enabled", False))
freeze_backbone = bool(model_cfg.get("freeze_backbone", False))
lora_cfg = model_cfg.get("lora", {})
lora_enabled = bool(lora_cfg.get("enabled", False))
if lora_enabled and model_name not in ("vit_s16_cls", "vit_l16_cls"):
raise ValueError("LoRA attention injection is currently supported only for vit_s16_cls and vit_l16_cls (timm VisionTransformer)")
if model_name == "resnet18_cls":
if use_aux_heads:
if aux_dims is None:
raise ValueError("aux_dims is required when aux_heads.enabled=true")
model = MultiHeadResNet18Classifier(
pretrained=pretrained,
num_classes=num_classes,
aux_dims=aux_dims,
dropout_p=dropout_p,
)
else:
model = resnet18(weights=_pretrained_weights(pretrained, ResNet18_Weights))
in_features = model.fc.in_features
model.fc = _make_head(in_features, num_classes, dropout_p)
elif model_name == "mobilenet_v3_small_cls":
model = mobilenet_v3_small(weights=_pretrained_weights(pretrained, MobileNet_V3_Small_Weights))
in_features = model.classifier[-1].in_features
model.classifier[-1] = _make_head(in_features, num_classes, dropout_p)
elif model_name == "efficientnet_b0_cls":
model = efficientnet_b0(weights=_pretrained_weights(pretrained, EfficientNet_B0_Weights))
in_features = model.classifier[-1].in_features
model.classifier[-1] = _make_head(in_features, num_classes, dropout_p)
elif model_name == "convnext_tiny_cls":
model = convnext_tiny(weights=_pretrained_weights(pretrained, ConvNeXt_Tiny_Weights))
in_features = model.classifier[-1].in_features
model.classifier[-1] = _make_head(in_features, num_classes, dropout_p)
elif model_name == "vit_b16_cls":
model = vit_b_16(weights=_pretrained_weights(pretrained, ViT_B_16_Weights))
in_features = model.heads.head.in_features
model.heads.head = _make_head(in_features, num_classes, dropout_p)
elif model_name in ("vit_s16_cls", "vit_l16_cls"):
timm_model_name = _resolve_timm_classifier_name(model_cfg)
if timm_model_name is None:
raise ValueError(f"timm_model_name could not be resolved for {model_name}")
model = timm.create_model(
timm_model_name,
pretrained=pretrained,
num_classes=num_classes,
drop_rate=dropout_p,
)
if lora_enabled:
lora_targets = tuple(lora_cfg.get("target_modules", ("qkv", "proj")))
injected_layers = apply_lora_to_timm_attention(
model,
rank=int(lora_cfg.get("rank", 8)),
alpha=float(lora_cfg.get("alpha", lora_cfg.get("rank", 8))),
dropout_p=float(lora_cfg.get("dropout", 0.0)),
target_modules=lora_targets,
last_n_blocks=int(lora_cfg["last_n_blocks"]) if lora_cfg.get("last_n_blocks") is not None else None,
block_indices=lora_cfg.get("block_indices"),
)
model._lora_injected_layers = injected_layers
target_block_indices = tuple(getattr(model, "_lora_target_block_indices", ()))
if freeze_backbone and target_block_indices:
try:
frozen_prefix_blocks = enable_lora_suffix_no_grad(model, target_block_indices)
model._lora_no_grad_prefix_count = frozen_prefix_blocks
except ValueError:
model._lora_no_grad_prefix_count = 0
elif model_name == "swin_t_cls":
model = swin_t(weights=_pretrained_weights(pretrained, Swin_T_Weights))
in_features = model.head.in_features
model.head = _make_head(in_features, num_classes, dropout_p)
else:
raise ValueError(
f"Unsupported classification model '{model_name}'. "
"Supported: resnet18_cls, mobilenet_v3_small_cls, efficientnet_b0_cls, "
"convnext_tiny_cls, vit_b16_cls, vit_s16_cls, swin_t_cls"
)
if freeze_backbone:
freeze_all_parameters_except_lora_and_heads(model)
return model
|