v4 update
This commit is contained in:
rpotter6298
@@ -0,0 +1,217 @@
|
||||
"""backbones — backbone registry and builder for v4 image towers."""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Callable, Dict, List
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torchvision import models, transforms
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
|
||||
REFUGELIKE_BACKBONE_PATH = _REPO_ROOT / "models/v2/refuge/refugelike_backbone.pt"
|
||||
REFUGE_DENSENET_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_densenet_backbone.pt"
|
||||
REFUGE_EFFICIENT_B0_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_efficient_b0_backbone.pt"
|
||||
REFUGE_EFFICIENT_B7_PATH = _REPO_ROOT / "models/refuge/classifier/refuge_efficient_b7_backbone.pt"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BackboneSpec:
|
||||
ctor: Callable
|
||||
weights_default: object
|
||||
strip: Callable[[nn.Module], tuple]
|
||||
blocks: Callable[[nn.Module], List[nn.Module]]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Strip helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _strip_efficientnet(m: models.EfficientNet):
|
||||
out_dim = m.classifier[1].in_features
|
||||
m.classifier = nn.Identity()
|
||||
return out_dim, m
|
||||
|
||||
def _strip_resnet(m: models.ResNet):
|
||||
out_dim = m.fc.in_features
|
||||
m.fc = nn.Identity()
|
||||
return out_dim, m
|
||||
|
||||
def _strip_densenet(m: models.DenseNet):
|
||||
out_dim = m.classifier.in_features
|
||||
m.classifier = nn.Identity()
|
||||
return out_dim, m
|
||||
|
||||
def _strip_vgg(m: models.VGG):
|
||||
out_dim = m.classifier[0].in_features
|
||||
m.classifier = nn.Identity()
|
||||
return out_dim, m
|
||||
|
||||
def _strip_mobilenet_v2(m: models.MobileNetV2):
|
||||
out_dim = m.classifier[1].in_features
|
||||
m.classifier = nn.Identity()
|
||||
return out_dim, m
|
||||
|
||||
def _strip_inception_v3(m: models.Inception3):
|
||||
out_dim = m.fc.in_features
|
||||
m.fc = nn.Identity()
|
||||
m.aux_logits = False
|
||||
m.AuxLogits = None
|
||||
return out_dim, m
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Block splitters for ratio-based freezing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _blocks_efficientnet(m: models.EfficientNet):
|
||||
return list(m.features)
|
||||
|
||||
def _blocks_resnet(m: models.ResNet):
|
||||
stem = nn.Sequential(m.conv1, m.bn1, m.relu, m.maxpool)
|
||||
return [stem, m.layer1, m.layer2, m.layer3, m.layer4]
|
||||
|
||||
def _blocks_densenet(m: models.DenseNet):
|
||||
f = m.features
|
||||
stem = nn.Sequential(f.conv0, f.norm0, f.relu0, f.pool0)
|
||||
return [stem, f.denseblock1, f.transition1, f.denseblock2, f.transition2,
|
||||
f.denseblock3, f.transition3, f.denseblock4, f.norm5]
|
||||
|
||||
def _blocks_vgg(m: models.VGG):
|
||||
stages, cur = [], []
|
||||
for mod in m.features:
|
||||
cur.append(mod)
|
||||
if isinstance(mod, nn.MaxPool2d):
|
||||
stages.append(nn.Sequential(*cur)); cur = []
|
||||
if cur:
|
||||
stages.append(nn.Sequential(*cur))
|
||||
return stages
|
||||
|
||||
def _blocks_mobilenet_v2(m: models.MobileNetV2):
|
||||
return list(m.features)
|
||||
|
||||
def _blocks_inception_v3(m: models.Inception3):
|
||||
return [child for name, child in m.named_children()
|
||||
if name not in ("fc", "AuxLogits")]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Registry
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
BACKBONES: Dict[str, BackboneSpec] = {
|
||||
"efficientnet_b0": BackboneSpec(
|
||||
ctor=models.efficientnet_b0,
|
||||
weights_default=models.EfficientNet_B0_Weights.DEFAULT,
|
||||
strip=_strip_efficientnet,
|
||||
blocks=_blocks_efficientnet,
|
||||
),
|
||||
"resnet50": BackboneSpec(
|
||||
ctor=models.resnet50,
|
||||
weights_default=models.ResNet50_Weights.DEFAULT,
|
||||
strip=_strip_resnet,
|
||||
blocks=_blocks_resnet,
|
||||
),
|
||||
"densenet121": BackboneSpec(
|
||||
ctor=models.densenet121,
|
||||
weights_default=models.DenseNet121_Weights.DEFAULT,
|
||||
strip=_strip_densenet,
|
||||
blocks=_blocks_densenet,
|
||||
),
|
||||
"vgg16": BackboneSpec(
|
||||
ctor=models.vgg16,
|
||||
weights_default=models.VGG16_Weights.DEFAULT,
|
||||
strip=_strip_vgg,
|
||||
blocks=_blocks_vgg,
|
||||
),
|
||||
"mobilenet_v2": BackboneSpec(
|
||||
ctor=models.mobilenet_v2,
|
||||
weights_default=models.MobileNet_V2_Weights.DEFAULT,
|
||||
strip=_strip_mobilenet_v2,
|
||||
blocks=_blocks_mobilenet_v2,
|
||||
),
|
||||
"inception_v3": BackboneSpec(
|
||||
ctor=models.inception_v3,
|
||||
weights_default=models.Inception_V3_Weights.DEFAULT,
|
||||
strip=_strip_inception_v3,
|
||||
blocks=_blocks_inception_v3,
|
||||
),
|
||||
"refugelike": BackboneSpec(
|
||||
ctor=models.resnet50,
|
||||
weights_default=None,
|
||||
strip=_strip_resnet,
|
||||
blocks=_blocks_resnet,
|
||||
),
|
||||
"refuge_densenet": BackboneSpec(
|
||||
ctor=models.densenet121,
|
||||
weights_default=None,
|
||||
strip=_strip_densenet,
|
||||
blocks=_blocks_densenet,
|
||||
),
|
||||
"refuge_efficient_b0": BackboneSpec(
|
||||
ctor=models.efficientnet_b0,
|
||||
weights_default=None,
|
||||
strip=_strip_efficientnet,
|
||||
blocks=_blocks_efficientnet,
|
||||
),
|
||||
"refuge_efficient_b7": BackboneSpec(
|
||||
ctor=models.efficientnet_b7,
|
||||
weights_default=None,
|
||||
strip=_strip_efficientnet,
|
||||
blocks=_blocks_efficientnet,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def list_names() -> List[str]:
|
||||
return list(BACKBONES.keys())
|
||||
|
||||
|
||||
def load_backbone_weights(key: str, model: nn.Module) -> None:
|
||||
paths = {
|
||||
"refugelike": REFUGELIKE_BACKBONE_PATH,
|
||||
"refuge_densenet": REFUGE_DENSENET_PATH,
|
||||
"refuge_efficient_b0": REFUGE_EFFICIENT_B0_PATH,
|
||||
"refuge_efficient_b7": REFUGE_EFFICIENT_B7_PATH,
|
||||
}
|
||||
path = paths.get(key)
|
||||
if path is None:
|
||||
return
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"Custom backbone weights not found at {path}. "
|
||||
"Export them via refuge_build.py --export-backbone first."
|
||||
)
|
||||
state = torch.load(path, map_location="cpu")
|
||||
model.load_state_dict(state, strict=False)
|
||||
|
||||
|
||||
def build_backbone(name: str, freeze_ratio: float = 0.0) -> tuple[nn.Module, int, list]:
|
||||
"""Instantiate a backbone, strip its classifier head, apply freeze ratio.
|
||||
|
||||
Returns (model, out_dim, blocks) where blocks is the ordered list of
|
||||
freezable units — callers use it to dynamically adjust freeze_ratio later.
|
||||
"""
|
||||
key = (name or "").lower()
|
||||
if key not in BACKBONES:
|
||||
raise ValueError(f"Unknown backbone '{name}'. Available: {list_names()}")
|
||||
|
||||
spec = BACKBONES[key]
|
||||
if spec.weights_default is not None:
|
||||
m = spec.ctor(weights=spec.weights_default)
|
||||
else:
|
||||
m = spec.ctor(weights=None)
|
||||
out_dim, m = spec.strip(m)
|
||||
load_backbone_weights(key, m)
|
||||
|
||||
blocks = spec.blocks(m)
|
||||
fr = max(0.0, min(1.0, float(freeze_ratio)))
|
||||
n_freeze = int(math.floor(len(blocks) * fr))
|
||||
for b in blocks[:n_freeze]:
|
||||
for p in b.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
return m, out_dim, blocks
|
||||
@@ -0,0 +1,126 @@
|
||||
"""se_block — Squeeze-and-Excitation channel gating."""
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class SEBlock(nn.Module):
|
||||
"""SE-style channel gating for vectors, feature maps, and sequences.
|
||||
|
||||
Input shapes:
|
||||
[N, C] (vector) — squeeze is identity
|
||||
[N, C, H, W] (image map) — squeeze over H, W
|
||||
[N, T, C] (sequence) — squeeze over T
|
||||
|
||||
Gate modes:
|
||||
residual (default): gate = 1 + tanh(MLP(s)) in (0, 2), identity at init
|
||||
plain: gate = sigmoid(MLP(s)) in (0, 1)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
reduction: int = 16,
|
||||
residual: bool = True,
|
||||
identity_init: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
hid = max(1, dim // max(1, reduction))
|
||||
self.fc1 = nn.Linear(dim, hid, bias=True)
|
||||
self.act = nn.ReLU(inplace=True)
|
||||
self.fc2 = nn.Linear(hid, dim, bias=True)
|
||||
self.residual = residual
|
||||
|
||||
if residual and identity_init:
|
||||
nn.init.zeros_(self.fc2.weight)
|
||||
nn.init.zeros_(self.fc2.bias)
|
||||
|
||||
def _squeeze(self, x: torch.Tensor) -> torch.Tensor:
|
||||
if x.dim() == 2:
|
||||
return x
|
||||
if x.dim() == 4:
|
||||
return x.mean(dim=(2, 3))
|
||||
if x.dim() == 3:
|
||||
return x.mean(dim=1)
|
||||
return x.view(x.size(0), -1)
|
||||
|
||||
def _broadcast(self, gate: torch.Tensor, like: torch.Tensor) -> torch.Tensor:
|
||||
if like.dim() == 2:
|
||||
return gate
|
||||
if like.dim() == 3:
|
||||
return gate.unsqueeze(1)
|
||||
if like.dim() == 4:
|
||||
return gate.unsqueeze(-1).unsqueeze(-1)
|
||||
return gate.view_as(like)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
s = self._squeeze(x)
|
||||
u = self.fc2(self.act(self.fc1(s)))
|
||||
gate = (1.0 + torch.tanh(u)) if self.residual else torch.sigmoid(u)
|
||||
return x * self._broadcast(gate, x), gate
|
||||
|
||||
|
||||
class SEGateLogger:
|
||||
"""Running stats over SE gate activations across batches."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
enabled: bool = True,
|
||||
track_channels: bool = False,
|
||||
dim: int | None = None,
|
||||
):
|
||||
self.enabled = enabled
|
||||
self.track_channels = track_channels
|
||||
self.dim = dim
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
self._n = 0
|
||||
self._sum = 0.0
|
||||
self._sum2 = 0.0
|
||||
self._lt02 = 0
|
||||
self._gt08 = 0
|
||||
self._ch_sum = None
|
||||
self._ch_count = 0
|
||||
if self.track_channels and self.dim is not None:
|
||||
self._ch_sum = torch.zeros(self.dim, dtype=torch.float32)
|
||||
|
||||
@torch.no_grad()
|
||||
def accumulate(self, gates: torch.Tensor) -> None:
|
||||
if not self.enabled:
|
||||
return
|
||||
if gates.dim() == 4:
|
||||
g = gates.mean(dim=(2, 3))
|
||||
elif gates.dim() == 3:
|
||||
g = gates.mean(dim=1)
|
||||
elif gates.dim() == 2:
|
||||
g = gates
|
||||
else:
|
||||
g = gates.view(gates.size(0), -1)
|
||||
g = g.detach()
|
||||
self._n += g.numel()
|
||||
self._sum += g.sum().item()
|
||||
self._sum2 += (g * g).sum().item()
|
||||
self._lt02 += (g < 0.2).sum().item()
|
||||
self._gt08 += (g > 0.8).sum().item()
|
||||
if self._ch_sum is not None:
|
||||
self._ch_sum += g.sum(dim=0).cpu()
|
||||
self._ch_count += g.size(0)
|
||||
|
||||
def get(self, reset: bool = True) -> dict | None:
|
||||
if self._n == 0:
|
||||
return None
|
||||
mean = self._sum / self._n
|
||||
var = max(0.0, self._sum2 / self._n - mean * mean)
|
||||
out = {
|
||||
"mean": mean,
|
||||
"std": var ** 0.5,
|
||||
"pct_lt_0.2": self._lt02 / self._n,
|
||||
"pct_gt_0.8": self._gt08 / self._n,
|
||||
}
|
||||
if self._ch_sum is not None and self._ch_count > 0:
|
||||
out["channel_mean"] = (self._ch_sum / float(self._ch_count)).tolist()
|
||||
if reset:
|
||||
self.reset()
|
||||
return out
|
||||
@@ -0,0 +1,66 @@
|
||||
"""transforms — image transform utilities for v4 towers."""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Tuple
|
||||
|
||||
from torchvision import transforms
|
||||
|
||||
from v4.classes.accessory.backbones import BACKBONES
|
||||
|
||||
IMAGENET_MEAN: Tuple[float, float, float] = (0.485, 0.456, 0.406)
|
||||
IMAGENET_STD: Tuple[float, float, float] = (0.229, 0.224, 0.225)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageTransformConfig:
|
||||
crop_size: int = 224
|
||||
resize_size: int = 256
|
||||
mean: Tuple[float, float, float] = IMAGENET_MEAN
|
||||
std: Tuple[float, float, float] = IMAGENET_STD
|
||||
augment: bool = True
|
||||
rotation_deg: int = 15
|
||||
color_jitter: Tuple[float, float, float, float] = (0.1, 0.1, 0.1, 0.05)
|
||||
hflip: bool = True
|
||||
vflip: bool = True
|
||||
|
||||
def build(self) -> transforms.Compose:
|
||||
ops = [
|
||||
transforms.Resize(self.resize_size),
|
||||
transforms.CenterCrop(self.crop_size),
|
||||
]
|
||||
if self.augment:
|
||||
if self.hflip:
|
||||
ops.append(transforms.RandomHorizontalFlip())
|
||||
if self.vflip:
|
||||
ops.append(transforms.RandomVerticalFlip())
|
||||
if self.rotation_deg:
|
||||
ops.append(transforms.RandomRotation(self.rotation_deg))
|
||||
if self.color_jitter:
|
||||
ops.append(transforms.ColorJitter(*self.color_jitter))
|
||||
ops += [
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=self.mean, std=self.std),
|
||||
]
|
||||
return transforms.Compose(ops)
|
||||
|
||||
|
||||
def backbone_transform_config(backbone_name: str, augment: bool = True) -> ImageTransformConfig:
|
||||
"""Build an ImageTransformConfig using the backbone's default normalisation stats."""
|
||||
key = (backbone_name or "").lower()
|
||||
if key not in BACKBONES:
|
||||
raise ValueError(f"Unknown backbone '{backbone_name}'.")
|
||||
spec = BACKBONES[key]
|
||||
mean = getattr(spec.weights_default, "meta", {}).get("mean", IMAGENET_MEAN)
|
||||
std = getattr(spec.weights_default, "meta", {}).get("std", IMAGENET_STD)
|
||||
crop = 299 if key == "inception_v3" else 224
|
||||
return ImageTransformConfig(crop_size=crop, mean=mean, std=std, augment=augment)
|
||||
|
||||
|
||||
def build_backbone_transform(backbone_name: str, augment: bool = True) -> transforms.Compose:
|
||||
return backbone_transform_config(backbone_name, augment=augment).build()
|
||||
|
||||
|
||||
def build_eval_transform(backbone_name: str) -> transforms.Compose:
|
||||
"""Deterministic eval transform — no augmentation, backbone-matched normalisation."""
|
||||
return build_backbone_transform(backbone_name, augment=False)
|
||||
Reference in New Issue
Block a user