Upload 3 files

mobile_net.py

@@ -0,0 +1,345 @@
+'''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)

mobilenetv3_large.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86adb5480e88d76e9631b25070a455beae6ec01d3eac41fed726ccbbf7fa92e4
+size 15743184

mobilenetv3_small.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d9976a3bd3c1292b09c9ba34ee79a5c30b187a4b10a564b93d5360ff03c0638
+size 6773776