Upload 2 files

model.py

@@ -0,0 +1,108 @@
+import json
+import torch
+from torch import nn, Tensor
+from loguru import logger
+from pathlib import Path
+
+from torchvision.transforms import ToTensor
+from torchvision.transforms.v2 import CenterCrop, Compose, Normalize
+
+
+import vits
+
+def _clean_moco_state_dict(state_dict: dict[str, Tensor], linear_keyword: str) -> dict[str, Tensor]:
+    """
+    Filters and renames keys from a MoCo state_dict.
+
+    It selects keys from the 'base_encoder', removes the given linear layer keyword,
+    and strips the 'module.base_encoder.' prefix.
+    """
+    for key in list(state_dict.keys()):
+        # Check if the key belongs to the base encoder's backbone
+        if key.startswith('module.base_encoder') and not key.startswith(f'module.base_encoder.{linear_keyword}'):
+            # Create a new key by stripping the prefix
+            new_key = key[len("module.base_encoder."):]
+            state_dict[new_key] = state_dict[key]
+
+        # Delete the old key (either renamed or unused)
+        del state_dict[key]
+
+    return state_dict
+
+def load_moco_encoder(
+        model: nn.Module,
+        weight_path: Path,
+        linear_keyword: str,
+) -> nn.Module:
+    """
+    Loads pre-trained MoCo weights into a given model instance (ResNet, ViT, etc.).
+
+    This function handles loading the checkpoint, cleaning the state dictionary keys,
+    and loading the weights into the model's backbone. It finishes by replacing
+    the model's linear head with an Identity layer to turn it into a feature extractor.
+
+    Args:
+        model: An instantiated PyTorch model (e.g., from timm or a custom module).
+        weight_path: Path to the .pth or .pt MoCo checkpoint file.
+        linear_keyword: The name of the final linear layer to exclude (e.g., 'fc' or 'head').
+
+    Returns:
+        The same model, with pre-trained backbone weights and the head replaced
+        by nn.Identity(), ready for feature extraction.
+    """
+    assert weight_path.exists(), f"Checkpoint not found at '{weight_path}'"
+    logger.info(f"=> Loading MoCo checkpoint from '{weight_path}'")
+
+    # Use weights_only=True for added security if the checkpoint doesn't contain pickled code
+    checkpoint = torch.load(weight_path, map_location="cpu", weights_only=True)
+
+    # Extract the state dictionary containing the model weights
+    state_dict = checkpoint["state_dict"]
+
+    # Clean the state_dict to match the model's architecture
+    cleaned_state_dict = _clean_moco_state_dict(state_dict, linear_keyword)
+
+    # Load the cleaned weights into the model
+    msg = model.load_state_dict(cleaned_state_dict, strict=False)
+    logger.info(msg)
+    logger.info("=> Successfully loaded pre-trained model backbone.")
+
+    # Replace the model's head to turn it into a feature extractor
+    if hasattr(model, linear_keyword):
+        setattr(model, linear_keyword, nn.Identity())
+        logger.info(f"=> Model's '{linear_keyword}' layer replaced with nn.Identity for feature extraction.")
+
+    return model
+
+def get_vit_feature_extractor(weight_path: Path, model_name: str = "vits8", img_size: int = 40) -> nn.Module:
+    """Creates a ViT feature extractor using the unified loader."""
+    # 1. Create the model architecture shell
+    vit_model = vits.__dict__[model_name](img_size=img_size, num_classes=0)
+
+    # 2. Use the unified function to load weights and prepare for feature extraction
+    feature_extractor = load_moco_encoder(
+        model=vit_model,
+        weight_path=weight_path,
+        linear_keyword='head'
+    )
+    return feature_extractor
+
+
+def prepare_transform(
+        stats_path,
+        size: int = 40,
+) -> Compose:
+    # Get normalisation stats
+    with open(stats_path, "r") as f:
+        norm_dict = json.load(f)
+    mean = norm_dict["mean"]
+    std = norm_dict["std"]
+
+    # Prepare transform
+    list_transform = [
+        ToTensor(),
+        Normalize(mean=mean, std=std),
+        CenterCrop(size=size),
+    ]
+    transform = Compose(list_transform)
+    return transform

vits.py

@@ -0,0 +1,177 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+import torch
+import torch.nn as nn
+from functools import partial, reduce
+from operator import mul
+
+from timm.layers import to_2tuple
+from timm.models.vision_transformer import VisionTransformer, _cfg
+from timm.layers import PatchEmbed
+
+__all__ = [
+    "vits4",
+    "vits8",
+    "vitb4",
+    "vitb8",
+]
+
+
+class VisionTransformerMoCo(VisionTransformer):
+    def __init__(self, stop_grad_conv1=False, **kwargs):
+        super().__init__(**kwargs)
+        # Use fixed 2D sin-cos position embedding
+        self.build_2d_sincos_position_embedding()
+
+        # weight initialization
+        for name, m in self.named_modules():
+            if isinstance(m, nn.Linear):
+                if "qkv" in name:
+                    # treat the weights of Q, K, V separately
+                    val = math.sqrt(
+                        6.0 / float(m.weight.shape[0] // 3 + m.weight.shape[1])
+                    )
+                    nn.init.uniform_(m.weight, -val, val)
+                else:
+                    nn.init.xavier_uniform_(m.weight)
+                nn.init.zeros_(m.bias)
+        nn.init.normal_(self.cls_token, std=1e-6)
+
+        if isinstance(self.patch_embed, PatchEmbed):
+            # xavier_uniform initialization
+            val = math.sqrt(
+                6.0
+                / float(
+                    3 * reduce(mul, self.patch_embed.patch_size, 1) + self.embed_dim
+                )
+            )
+            nn.init.uniform_(self.patch_embed.proj.weight, -val, val)
+            nn.init.zeros_(self.patch_embed.proj.bias)
+
+            if stop_grad_conv1:
+                self.patch_embed.proj.weight.requires_grad = False
+                self.patch_embed.proj.bias.requires_grad = False
+
+    def build_2d_sincos_position_embedding(self, temperature=10000.0):
+        h, w = self.patch_embed.grid_size
+        grid_w = torch.arange(w, dtype=torch.float32)
+        grid_h = torch.arange(h, dtype=torch.float32)
+        grid_w, grid_h = torch.meshgrid(grid_w, grid_h)
+        assert (
+            self.embed_dim % 4 == 0
+        ), "Embed dimension must be divisible by 4 for 2D sin-cos position embedding"
+        pos_dim = self.embed_dim // 4
+        omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim
+        omega = 1.0 / (temperature**omega)
+        out_w = torch.einsum("m,d->md", [grid_w.flatten(), omega])
+        out_h = torch.einsum("m,d->md", [grid_h.flatten(), omega])
+        pos_emb = torch.cat(
+            [torch.sin(out_w), torch.cos(out_w), torch.sin(out_h), torch.cos(out_h)],
+            dim=1,
+        )[None, :, :]
+
+        pe_token = torch.zeros([1, self.num_prefix_tokens, self.embed_dim], dtype=torch.float32)
+        self.pos_embed = nn.Parameter(torch.cat([pe_token, pos_emb], dim=1))
+        self.pos_embed.requires_grad = False
+
+
+class ConvStem(nn.Module):
+    """
+    ConvStem, from Early Convolutions Help Transformers See Better, Tete et al. https://arxiv.org/abs/2106.14881
+    """
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        norm_layer=None,
+        flatten=True,
+    ):
+        super().__init__()
+
+        assert patch_size == 16, "ConvStem only supports patch size of 16"
+        assert embed_dim % 8 == 0, "Embed dimension must be divisible by 8 for ConvStem"
+
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        # build stem, similar to the design in https://arxiv.org/abs/2106.14881
+        stem = []
+        input_dim, output_dim = 3, embed_dim // 8
+        for l in range(4):
+            stem.append(
+                nn.Conv2d(
+                    input_dim,
+                    output_dim,
+                    kernel_size=3,
+                    stride=2,
+                    padding=1,
+                    bias=False,
+                )
+            )
+            stem.append(nn.BatchNorm2d(output_dim))
+            stem.append(nn.ReLU(inplace=True))
+            input_dim = output_dim
+            output_dim *= 2
+        stem.append(nn.Conv2d(input_dim, embed_dim, kernel_size=1))
+        self.proj = nn.Sequential(*stem)
+
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert (
+            H == self.img_size[0] and W == self.img_size[1]
+        ), f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        x = self.norm(x)
+        return x
+
+
+def vits(patch_size: int, **kwargs):
+    model = VisionTransformerMoCo(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs,
+    )
+    model.default_cfg = _cfg()
+    return model
+
+vits4 = partial(vits, patch_size=4)
+vits8 = partial(vits, patch_size=8)
+
+def vitb(patch_size: int, **kwargs):
+    model = VisionTransformerMoCo(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs,
+    )
+    model.default_cfg = _cfg()
+    return model
+
+vitb4 = partial(vitb, patch_size=4)
+vitb8 = partial(vitb, patch_size=8)