| from __future__ import annotations |
|
|
| import json |
| from functools import partial |
| from pathlib import Path |
| from typing import Callable |
|
|
|
|
| import torch |
| from timm.layers import to_2tuple |
| from timm.models.vision_transformer import VisionTransformer |
| from torch import Tensor, nn |
| from torchvision.transforms import ToTensor |
| from torchvision.transforms.v2 import Compose, Normalize, Resize |
|
|
|
|
| TARGET_CELL_SIZE = 40 |
| LEMON_MOCO_ARCHITECTURES = ("vits8", "vitb8") |
|
|
|
|
| def build_lemon_moco_backbone(arch: str) -> nn.Module: |
| if arch not in LEMON_MOCO_ARCHITECTURES: |
| raise ValueError( |
| f"Unknown LEMON MoCo architecture: {arch}. " |
| f"Expected one of {LEMON_MOCO_ARCHITECTURES}." |
| ) |
| patch_size = 8 |
| embed_dim = 384 if arch == "vits8" else 768 |
| model = VisionTransformer( |
| img_size=TARGET_CELL_SIZE, |
| patch_size=patch_size, |
| embed_dim=embed_dim, |
| depth=12, |
| num_heads=12, |
| mlp_ratio=4, |
| qkv_bias=True, |
| norm_layer=partial(nn.LayerNorm, eps=1e-6), |
| num_classes=0, |
| ) |
| return model |
|
|
|
|
|
|
| def load_lemon_moco_normalization() -> tuple[list[float], list[float]]: |
| stats_path = Path(".") / "mean_std.json" |
| if not stats_path.is_file(): |
| raise FileNotFoundError(f"LEMON MoCo normalization stats are missing: {stats_path}") |
| with stats_path.open("r", encoding="utf-8") as stats_file: |
| stats = json.load(stats_file) |
| mean = [float(value) for value in stats["mean"]] |
| std = [float(value) for value in stats["std"]] |
| return mean, std |
|
|
|
|
| def build_lemon_moco_transform() -> Compose: |
| mean, std = load_lemon_moco_normalization() |
| transform = Compose( |
| [ |
| ToTensor(), |
| Resize(size=to_2tuple(TARGET_CELL_SIZE)), |
| Normalize(mean=mean, std=std), |
| ] |
| ) |
| return transform |
|
|
|
|
| def extract_backbone_state_dict(checkpoint: dict) -> dict[str, Tensor]: |
| if "state_dict" not in checkpoint: |
| raise KeyError("LEMON MoCo checkpoint does not contain a 'state_dict' entry.") |
| backbone_state_dict = {} |
| for raw_key, value in checkpoint["state_dict"].items(): |
| key = raw_key.removeprefix("module.") |
| if key.startswith("base_encoder.") and not key.startswith( |
| ("base_encoder.head.", "base_encoder.fc.") |
| ): |
| backbone_state_dict[key.removeprefix("base_encoder.")] = value |
| if not backbone_state_dict: |
| raise ValueError("LEMON MoCo checkpoint did not contain base_encoder backbone weights.") |
| return backbone_state_dict |
|
|
|
|
| def load_lemon_moco_model(arch: str = "vitb8") -> tuple[nn.Module, torch.dtype, Callable]: |
| checkpoint_path = Path(".") / f"lemon_{arch}.pth.tar" |
| if not checkpoint_path.is_file(): |
| raise FileNotFoundError(f"LEMON MoCo checkpoint does not exist: {checkpoint_path}") |
| transform = build_lemon_moco_transform() |
| model = build_lemon_moco_backbone(arch) |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) |
| backbone_state_dict = extract_backbone_state_dict(checkpoint) |
| incompatible_keys = model.load_state_dict(backbone_state_dict, strict=False) |
| if incompatible_keys.missing_keys: |
| raise RuntimeError( |
| f"Missing keys when loading LEMON MoCo backbone: {incompatible_keys.missing_keys}" |
| ) |
| if incompatible_keys.unexpected_keys: |
| raise RuntimeError( |
| f"Unexpected keys when loading LEMON MoCo backbone: {incompatible_keys.unexpected_keys}" |
| ) |
| return model, torch.float16, transform |
|
|
