Upload projection head checkpoint, inference script, and asset

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+asset/umbreon2.jpg filter=lfs diff=lfs merge=lfs -text

PokeCon_head.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7b8a1e98908db06bdd9efd6566812517e6946e8feebaa8b1915aa88dfc996a2f
+size 9468377

inference.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoImageProcessor, AutoModel
+from transformers.image_utils import load_image
+
+MODEL_DIR = "your/path/to/facebook/dinov3-vith16plus-pretrain-lvd1689m"
+CKPT_PATH = "your/path/to/PokeCon_head.pt"
+
+IMG1 = "your_image1.png"
+IMG2 = "your_image2.png"
+
+DTYPE = torch.bfloat16  
+
+class ProjectionHead(nn.Module):
+    def __init__(self, in_dim: int, hidden_dims: list[int], dropout: float = 0.05, use_layernorm: bool = True):
+        super().__init__()
+        layers: list[nn.Module] = []
+        prev = in_dim
+        for i, d in enumerate(hidden_dims):
+            layers.append(nn.Linear(prev, d))
+            is_last = (i == len(hidden_dims) - 1)
+            if not is_last:
+                if use_layernorm:
+                    layers.append(nn.LayerNorm(d))
+                layers.append(nn.GELU())
+                if dropout > 0:
+                    layers.append(nn.Dropout(dropout))
+            prev = d
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+def main() -> None:
+    image1 = load_image(IMG1)
+    image2 = load_image(IMG2)
+
+    processor = AutoImageProcessor.from_pretrained(MODEL_DIR, local_files_only=True)
+    backbone = AutoModel.from_pretrained(
+        MODEL_DIR,
+        local_files_only=True,
+        torch_dtype=DTYPE,
+        device_map="auto",
+    ).eval()
+
+    ckpt = torch.load(CKPT_PATH, map_location="cpu")
+    if not isinstance(ckpt, dict) or "head_state_dict" not in ckpt:
+        raise RuntimeError("Expected a checkpoint dict with key: 'head_state_dict'")
+
+    cfg = ckpt.get("config", {})
+    proj_dims = list(cfg.get("proj_dims", [512, 256]))
+    dropout = float(cfg.get("dropout", 0.05))
+    use_layernorm = bool(cfg.get("use_layernorm", True))
+
+    head = ProjectionHead(
+        in_dim=backbone.config.hidden_size,
+        hidden_dims=proj_dims,
+        dropout=dropout,
+        use_layernorm=use_layernorm,
+    ).to(device=backbone.device, dtype=DTYPE)
+    head.load_state_dict(ckpt["head_state_dict"], strict=True)
+    head.eval()
+
+    inputs = processor(images=[image1, image2], return_tensors="pt").to(backbone.device)
+
+    with torch.inference_mode():
+        out = backbone(pixel_values=inputs["pixel_values"].to(backbone.dtype))
+        pooled = out.pooler_output  
+        z = head(pooled.to(DTYPE))
+        z = F.normalize(z, dim=-1)
+        cos = (z[0] * z[1]).sum().item()
+
+    print(f"Cosine similarity: {cos:.6f}")
+
+
+if __name__ == "__main__":
+    main()