mobile_net.py

'''MobileNetV3 feature extractors (Small / Large), refactored from
https://github.com/xiaolai-sqlai/mobilenetv3.

Changes vs. the original classification model:
  - classification head (linear4) removed, backbone only;
  - unified naming: bn1/bn2/bn3 -> norm1/norm2/norm3, linear3 -> proj,
    Block.se.se.* -> Block.se.features.*;
  - `MobileNetV3` takes a `backend` ('small' | 'large'); its classmethod
    `from_pretrained` infers the backend from a checkpoint file automatically.
'''
import torch
import torch.nn as nn
from torch.nn import init

from safetensors.torch import load_file, save_file

# (kernel, in_ch, expand_ch, out_ch, is_relu, use_se, stride)
_BACKEND_BLOCKS = {
    'small': [
        (3, 16, 16, 16, True, True, 2),
        (3, 16, 72, 24, True, False, 2),
        (3, 24, 88, 24, True, False, 1),
        (5, 24, 96, 40, False, True, 2),
        (5, 40, 240, 40, False, True, 1),
        (5, 40, 240, 40, False, True, 1),
        (5, 40, 120, 48, False, True, 1),
        (5, 48, 144, 48, False, True, 1),
        (5, 48, 288, 96, False, True, 2),
        (5, 96, 576, 96, False, True, 1),
        (5, 96, 576, 96, False, True, 1),
    ],
    'large': [
        (3, 16, 16, 16, True, False, 1),
        (3, 16, 64, 24, True, False, 2),
        (3, 24, 72, 24, True, False, 1),
        (5, 24, 72, 40, True, True, 2),
        (5, 40, 120, 40, True, True, 1),
        (5, 40, 120, 40, True, True, 1),
        (3, 40, 240, 80, False, False, 2),
        (3, 80, 200, 80, False, False, 1),
        (3, 80, 184, 80, False, False, 1),
        (3, 80, 184, 80, False, False, 1),
        (3, 80, 480, 112, False, True, 1),
        (3, 112, 672, 112, False, True, 1),
        (5, 112, 672, 160, False, True, 2),
        (5, 160, 672, 160, False, True, 1),
        (5, 160, 960, 160, False, True, 1),
    ],
}

# (head_in_ch, head_out_ch) fed to conv2 / proj
_BACKEND_HEAD = {
    'small': (96, 576),
    'large': (160, 960),
}

# conv2.weight (out, in, 1, 1) used to tell Small from Large
_BACKEND_CONV2_SHAPE = {
    'small': (576, 96),
    'large': (960, 160),
}

# Number of bneck blocks per backend (fallback signature when conv2 is absent)
_BACKEND_NBLOCKS = {
    'small': 11,
    'large': 15,
}


class SEModule(nn.Module):
    '''Squeeze-and-excitation block (same layout as upstream, feature extractor only).'''

    def __init__(self, in_size, reduction=4):
        super(SEModule, self).__init__()
        expand_size = max(in_size // reduction, 8)

        self.features = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_size, expand_size, kernel_size=1, bias=False),
            nn.BatchNorm2d(expand_size),
            nn.ReLU(inplace=True),
            nn.Conv2d(expand_size, in_size, kernel_size=1, bias=False),
            nn.Hardsigmoid(),
        )

    def forward(self, x):
        return x * self.features(x)


class Block(nn.Module):
    '''expand + depthwise + pointwise.'''

    def __init__(self, kernel_size, in_size, expand_size, out_size, act, se, stride):
        super(Block, self).__init__()
        self.stride = stride

        self.conv1 = nn.Conv2d(in_size, expand_size, kernel_size=1, bias=False)
        self.norm1 = nn.BatchNorm2d(expand_size)
        self.act1 = act(inplace=True)

        self.conv2 = nn.Conv2d(
            expand_size, expand_size, kernel_size=kernel_size, stride=stride,
            padding=kernel_size // 2, groups=expand_size, bias=False,
        )
        self.norm2 = nn.BatchNorm2d(expand_size)
        self.act2 = act(inplace=True)

        self.se = SEModule(expand_size) if se else nn.Identity()

        self.conv3 = nn.Conv2d(expand_size, out_size, kernel_size=1, bias=False)
        self.norm3 = nn.BatchNorm2d(out_size)
        self.act3 = act(inplace=True)

        self.skip = None
        if stride == 1 and in_size != out_size:
            self.skip = nn.Sequential(
                nn.Conv2d(in_size, out_size, kernel_size=1, bias=False),
                nn.BatchNorm2d(out_size),
            )
        if stride == 2 and in_size != out_size:
            self.skip = nn.Sequential(
                nn.Conv2d(in_channels=in_size, out_channels=in_size, kernel_size=3,
                          groups=in_size, stride=2, padding=1, bias=False),
                nn.BatchNorm2d(in_size),
                nn.Conv2d(in_size, out_size, kernel_size=1, bias=True),
                nn.BatchNorm2d(out_size),
            )
        if stride == 2 and in_size == out_size:
            self.skip = nn.Sequential(
                nn.Conv2d(in_channels=in_size, out_channels=out_size, kernel_size=3,
                          groups=in_size, stride=2, padding=1, bias=False),
                nn.BatchNorm2d(out_size),
            )

    def forward(self, x):
        skip = x

        out = self.act1(self.norm1(self.conv1(x)))
        out = self.act2(self.norm2(self.conv2(out)))
        out = self.se(out)
        out = self.norm3(self.conv3(out))

        if self.skip is not None:
            skip = self.skip(skip)
        return self.act3(out + skip)


def _read_tensors(path: str):
    '''Read a checkpoint into a {key: Tensor} dict (any tensors only).

    Supports .safetensors and .pth/.pt. State-dict wrappers
    ({'state_dict': ...} / {'model': ...}) and a DataParallel 'module.' prefix
    are handled transparently. No strictness checks here.
    '''
    if path.endswith('.safetensors'):
        raw = load_file(path, device='cpu')
    elif path.endswith(('.pth', '.pt')):
        raw = torch.load(path, map_location='cpu')
        if isinstance(raw, dict):
            for wrapper in ('state_dict', 'model'):
                sub = raw.get(wrapper)
                if isinstance(sub, dict):
                    raw = sub
                    break
        if not isinstance(raw, dict):
            raise RuntimeError(f'{path} is not a valid PyTorch weight file (expected a dict)')
    else:
        raise ValueError(
            f'unsupported weight format (only .safetensors / .pth / .pt): {path!r}')

    tensors = {}
    for key, val in raw.items():
        if not isinstance(val, torch.Tensor):
            continue  # skip non-weight entries such as epoch / optimizer
        if key.startswith('module.'):
            key = key[len('module.'):]  # strip DataParallel prefix
        tensors[key] = val
    return tensors


def detect_backend(path: str) -> str:
    '''Return 'small' or 'large' for the backend stored in a checkpoint file.'''
    tensors = _read_tensors(path)

    conv2 = tensors.get('conv2.weight')
    if conv2 is not None:
        shape = tuple(conv2.shape[:2])
        for name, expected in _BACKEND_CONV2_SHAPE.items():
            if shape == expected:
                return name
        raise ValueError(
            f'cannot tell Small from Large: conv2.weight shape {shape} matches neither '
            f'{_BACKEND_CONV2_SHAPE}')

    n_blocks = max(
        (int(key.split('.')[1]) for key in tensors if key.startswith('bneck.') and key.split('.')[1].isdigit()),
        default=-1,
    ) + 1
    for name, expected in _BACKEND_NBLOCKS.items():
        if n_blocks == expected:
            return name
    raise ValueError(
        f'cannot tell Small from Large: {n_blocks} bneck blocks match neither '
        f'{_BACKEND_NBLOCKS}')


class MobileNetV3(nn.Module):
    '''MobileNetV3 feature extractor. Outputs a 1280-dim feature vector per image.

    Args:
        backend: 'small' or 'large'. Defaults to 'small' for a bare instance;
            prefer `MobileNetV3.from_pretrained(path)` to pick it automatically.
        act: activation used by the hard-swish blocks (default nn.Hardswish).
    '''

    backend = None

    def __init__(self, backend: str | None = None, act=nn.Hardswish):
        super(MobileNetV3, self).__init__()
        if backend is None:
            backend = 'small' if self.backend is None else self.backend
        if backend not in _BACKEND_BLOCKS:
            raise ValueError(f'unknown backend {backend!r}; choose from {list(_BACKEND_BLOCKS)}')
        self.backend = backend

        head_in, head_out = _BACKEND_HEAD[backend]

        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1, bias=False)
        self.norm1 = nn.BatchNorm2d(16)
        self.act1 = act(inplace=True)

        def act_for(is_relu):
            return nn.ReLU if is_relu else act

        self.bneck = nn.Sequential(*[
            Block(k, i, e, o, act_for(relu), se, s)
            for (k, i, e, o, relu, se, s) in _BACKEND_BLOCKS[backend]
        ])

        self.conv2 = nn.Conv2d(head_in, head_out, kernel_size=1, stride=1, padding=0, bias=False)
        self.norm2 = nn.BatchNorm2d(head_out)
        self.act2 = act(inplace=True)
        self.gap = nn.AdaptiveAvgPool2d(1)

        self.proj = nn.Linear(head_out, 1280, bias=False)
        self.norm3 = nn.BatchNorm1d(1280)
        self.act3 = act(inplace=True)
        self.drop = nn.Dropout(0.2)

        self.init_params()

    def init_params(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        out = self.act1(self.norm1(self.conv1(x)))
        out = self.bneck(out)

        out = self.act2(self.norm2(self.conv2(out)))
        out = self.gap(out).flatten(1)
        out = self.drop(self.act3(self.norm3(self.proj(out))))

        return out

    def save_pretrained(self, path: str):
        '''Save the current weights by extension: safetensors or torch .pth/.pt.'''
        sd = self.state_dict()
        if path.endswith('.safetensors'):
            save_file(sd, path)
        elif path.endswith(('.pth', '.pt')):
            torch.save(sd, path)
        else:
            raise ValueError(
                f'unsupported weight format (only .safetensors / .pth / .pt): {path!r}')
        return self

    def load_pretrained(self, path: str):
        '''Load weights by extension (.safetensors / .pth / .pt).

        Strictly requires the current naming: keys must match this model exactly
        (no extra, none missing, per-tensor shapes equal). No legacy fallback.
        '''
        tensors = _read_tensors(path)

        ref = self.state_dict()
        extra = sorted(k for k in tensors if k not in ref)
        missing = sorted(k for k in ref if k not in tensors)
        if extra or missing:
            raise RuntimeError(
                f'weights do not match this model ({self.backend}): {len(extra)} extra / '
                f'{len(missing)} missing -> extra {extra[:5]}..., missing {missing[:5]}...')

        for k, v in tensors.items():
            want = tuple(ref[k].shape)
            if tuple(v.shape) != want:
                raise RuntimeError(f'shape mismatch for {k}: weights {tuple(v.shape)} vs model {want}')
            if v.dtype != ref[k].dtype:
                tensors[k] = v.to(ref[k].dtype)
        self.load_state_dict(tensors, strict=True)
        return self

    @classmethod
    def from_pretrained(cls, path: str) -> 'MobileNetV3':
        '''Infer the backend ('small'/'large') from the checkpoint and load it.

        Calling it on a pinned subclass raises if that subclass disagrees with
        the backend detected in the file.
        '''
        backend = detect_backend(path)
        pinned = cls.backend
        if pinned is not None and pinned != backend:
            raise ValueError(
                f'checkpoint at {path!r} is a {backend} model, but {cls.__name__} '
                f'is pinned to {pinned!r}')
        model = cls(backend=backend) if pinned is None else cls()
        return model.load_pretrained(path)


class MobileNetV3_Small(MobileNetV3):
    '''MobileNetV3-Small feature extractor (explicit backend, no auto-detection).'''

    backend = 'small'

    def __init__(self, act=nn.Hardswish):
        super(MobileNetV3_Small, self).__init__(backend='small', act=act)


class MobileNetV3_Large(MobileNetV3):
    '''MobileNetV3-Large feature extractor (explicit backend, no auto-detection).'''

    backend = 'large'

    def __init__(self, act=nn.Hardswish):
        super(MobileNetV3_Large, self).__init__(backend='large', act=act)