{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceType":"competition","sourceId":126777,"databundleVersionId":15314950}],"dockerImageVersionId":31286,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Experiment 13: ArcFace Hyperparameter Sweep on DINOv2 Embeddings\n\n## Research Question\n\nIs the default ArcFace configuration (margin=0.5, scale=64, embdim=256) optimal for\njaguar re-ID on frozen DINOv2-ViT-L/14 embeddings, or does a 3-stage grid sweep of\nmargin, scale, embedding dimension, hidden dimension, and dropout push mAP beyond the\nExp 4 baseline of 0.8447?\n\n## Overview\n\n1. Load cached DINOv2-ViT-L-14 embeddings — no image loading at any stage\n2. Stage 1: Sweep margin (0.2, 0.3, 0.4, 0.5, 0.6, 0.7) × scale (48, 64, 80, 128)\n3. Stage 2: Sweep embedding dim (128, 256, 512, 1024) using best margin+scale from Stage 1\n4. Stage 3: Sweep hidden dim (512, 1024) and dropout (0.1, 0.3, 0.5) using best settings from Stages 1 and 2\n5. Pick best overall configuration and generate submission\n\n## Controlled Variables\n\nBackbone: DINOv2-ViT-L-14, frozen throughout (Exp 4 cache reused).\nOptimizer: AdamW (lr=1e-4, weight_decay=1e-4).\nScheduler: ReduceLROnPlateau (factor=0.5, patience=10).\nEpochs: 50, batch size: 32, seed: 42.\nBaseline: mAP 0.8447 (margin=0.5, scale=64, dim=256, hidden=512, dropout=0.3)","metadata":{}},{"cell_type":"markdown","source":"## 1. Setup","metadata":{}},{"cell_type":"code","source":"import os, math, time, copy\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nimport timm\nfrom torchvision import transforms\nfrom PIL import Image\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom collections import defaultdict\nfrom tqdm.notebook import tqdm\nimport matplotlib.pyplot as plt\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import LabelEncoder\nimport wandb\nfrom kaggle_secrets import UserSecretsClient\n\nuser_secrets = UserSecretsClient()\nos.environ[\"HF_TOKEN\"]      = user_secrets.get_secret(\"hf_api\")\nos.environ[\"WANDB_API_KEY\"] = user_secrets.get_secret(\"wandb_api\")\n\nRANDOM_SEED = 42\n\ndef set_seed(seed):\n    import random\n    random.seed(seed); np.random.seed(seed)\n    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark     = False\n\nset_seed(RANDOM_SEED)\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Device: {device}  |  PyTorch: {torch.__version__}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:09:38.572439Z","iopub.execute_input":"2026-03-12T18:09:38.572822Z","iopub.status.idle":"2026-03-12T18:09:51.503232Z","shell.execute_reply.started":"2026-03-12T18:09:38.572787Z","shell.execute_reply":"2026-03-12T18:09:51.502589Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"RANDOM_SEED = 42\n\n# Baseline from Exp \nEXP4_BASELINE_MAP = 0.8447\n\n# Stage 1: Margin × Scale sweep\nSTAGE1_MARGINS = [0.2, 0.3, 0.4, 0.5, 0.6, 0.7]\nSTAGE1_SCALES  = [48.0, 64.0, 80.0, 128.0]\n\n# Stage 2: Embedding dimension (uses best margin+scale from Stage 1)\nSTAGE2_EMB_DIMS = [128, 256, 512, 1024]\n\n# Stage 3: Hidden dim + Dropout (uses best from Stages 1+2)\nSTAGE3_CONFIGS = [\n    {\"hidden_dim\": 512,  \"dropout\": 0.1},\n    {\"hidden_dim\": 512,  \"dropout\": 0.3},   # ← Exp 4 default\n    {\"hidden_dim\": 512,  \"dropout\": 0.5},\n    {\"hidden_dim\": 1024, \"dropout\": 0.1},\n    {\"hidden_dim\": 1024, \"dropout\": 0.3},\n    {\"hidden_dim\": 1024, \"dropout\": 0.5},\n]\n\nconfig = {\n    \"data_dir\":   Path(\"/kaggle/input/competitions/jaguar-re-id\"),\n    \"cache_dir\":  Path(\"/kaggle/working/embeddings\"),\n\n    # Backbone (frozen — embeddings reused from Exp 4 / Exp 9 cache)\n    \"dinov2_model\":      \"vit_large_patch14_dinov2.lvd142m\",\n    \"dinov2_input_size\": 518,\n\n    # Fixed training hyperparameters (identical to Exp 4)\n    \"batch_size\":    32,\n    \"learning_rate\": 1e-4,\n    \"weight_decay\":  1e-4,\n    \"num_epochs\":    50,\n    \"patience\":      10,\n    \"val_split\":     0.2,\n    \"seed\":          RANDOM_SEED,\n}\n\nconfig[\"cache_dir\"].mkdir(parents=True, exist_ok=True)\n\ntotal_s1 = len(STAGE1_MARGINS) * len(STAGE1_SCALES)\ntotal_s2 = len(STAGE2_EMB_DIMS)\ntotal_s3 = len(STAGE3_CONFIGS)\nprint(f\"Exp 4 baseline: mAP={EXP4_BASELINE_MAP} (margin=0.5, scale=64, dim=256, hidden=512, dropout=0.3)\")\nprint(f\"Stage 1 configs: {total_s1}  (margin × scale)\")\nprint(f\"Stage 2 configs: {total_s2}  (embedding dim)\")\nprint(f\"Stage 3 configs: {total_s3}  (hidden dim + dropout)\")\nprint(f\"All stages use cached DINOv2 embeddings — no image loading\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:09:51.504469Z","iopub.execute_input":"2026-03-12T18:09:51.504900Z","iopub.status.idle":"2026-03-12T18:09:51.512634Z","shell.execute_reply.started":"2026-03-12T18:09:51.504873Z","shell.execute_reply":"2026-03-12T18:09:51.511880Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"wandb.login()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:09:51.513497Z","iopub.execute_input":"2026-03-12T18:09:51.513774Z","iopub.status.idle":"2026-03-12T18:09:59.544776Z","shell.execute_reply.started":"2026-03-12T18:09:51.513752Z","shell.execute_reply":"2026-03-12T18:09:59.544021Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"wandb.init(\n    project=os.getenv(\"WANDB_PROJECT\", \"jaguar-reid-iota\"),\n    config={\n        \"experiment\":        \"arcface-sweep-dinov2\",\n        \"backbone\":          \"DINOv2-ViT-L-14 (frozen, Exp4 cache)\",\n        \"loss\":              \"CrossEntropy + ArcFace\",\n        \"stage1_margins\":    STAGE1_MARGINS,\n        \"stage1_scales\":     STAGE1_SCALES,\n        \"stage2_emb_dims\":   STAGE2_EMB_DIMS,\n        \"stage3_configs\":    STAGE3_CONFIGS,\n        \"exp4_baseline\":     EXP4_BASELINE_MAP,\n        \"lr\":                config[\"learning_rate\"],\n        \"batch_size\":        config[\"batch_size\"],\n        \"num_epochs\":        config[\"num_epochs\"],\n        \"seed\":              RANDOM_SEED,\n    },\n    name=\"arcface-sweep-dinov2\",\n)\nprint(\"W&B run: arcface-sweep-dinov2\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:09:59.545679Z","iopub.execute_input":"2026-03-12T18:09:59.546152Z","iopub.status.idle":"2026-03-12T18:10:08.583537Z","shell.execute_reply.started":"2026-03-12T18:09:59.546127Z","shell.execute_reply":"2026-03-12T18:10:08.582902Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. Load Data and Cached DINOv2 Embeddings","metadata":{}},{"cell_type":"code","source":"train_df = pd.read_csv(config[\"data_dir\"] / \"train.csv\")\nlabel_encoder = LabelEncoder()\ntrain_df['label_encoded'] = label_encoder.fit_transform(train_df['ground_truth'])\nnum_classes = len(label_encoder.classes_)\n\ntrain_data, val_data = train_test_split(\n    train_df, test_size=config[\"val_split\"],\n    random_state=config[\"seed\"], stratify=train_df['ground_truth'])\n\ntrain_labels_enc = train_data['label_encoded'].values\nval_labels_enc   = label_encoder.transform(val_data['ground_truth'].values)\nval_labels_str   = val_data['ground_truth'].values\n\nprint(f\"Train: {len(train_data)}  |  Val: {len(val_data)}  |  Classes: {num_classes}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:10:08.585433Z","iopub.execute_input":"2026-03-12T18:10:08.585759Z","iopub.status.idle":"2026-03-12T18:10:08.619054Z","shell.execute_reply.started":"2026-03-12T18:10:08.585733Z","shell.execute_reply":"2026-03-12T18:10:08.618508Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"d_tr_cache = config[\"cache_dir\"] / \"dinov2_train_embeddings.npz\"\nd_va_cache = config[\"cache_dir\"] / \"dinov2_val_embeddings.npz\"\n\nif d_tr_cache.exists() and d_va_cache.exists():\n    train_bb = np.load(d_tr_cache)[\"embeddings\"]\n    val_bb   = np.load(d_va_cache)[\"embeddings\"]\n    print(f\"Loaded cached DINOv2 embeddings ✓\")\n    print(f\"  Train: {train_bb.shape}  Val: {val_bb.shape}\")\n    dinov2_dim = train_bb.shape[1]\nelse:\n    # Extract from scratch if cache missing\n    print(\"Cache not found — extracting DINOv2 embeddings...\")\n    train_paths = [config[\"data_dir\"] / \"train/train\" / fn\n                   for fn in train_data[\"filename\"].astype(str)]\n    val_paths   = [config[\"data_dir\"] / \"train/train\" / fn\n                   for fn in val_data[\"filename\"].astype(str)]\n\n    backbone = timm.create_model(config[\"dinov2_model\"], pretrained=True,\n                                  num_classes=0, img_size=config[\"dinov2_input_size\"])\n    backbone.eval()\n    for p in backbone.parameters(): p.requires_grad = False\n    backbone.to(device)\n    with torch.no_grad():\n        dinov2_dim = backbone(torch.randn(1,3,518,518).to(device)).shape[1]\n\n    preprocess = transforms.Compose([\n        transforms.Resize((518, 518)), transforms.ToTensor(),\n        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),\n    ])\n\n    @torch.no_grad()\n    def extract(paths, desc):\n        out = []\n        for i in tqdm(range(0, len(paths), 32), desc=desc):\n            tensors = []\n            for p in paths[i:i+32]:\n                try:    tensors.append(preprocess(Image.open(p).convert(\"RGB\")))\n                except: tensors.append(torch.zeros(3,518,518))\n            out.append(backbone(torch.stack(tensors).to(device)).cpu().numpy())\n        return np.vstack(out)\n\n    train_bb = extract(train_paths, \"DINOv2 train\")\n    val_bb   = extract(val_paths,   \"DINOv2 val\")\n    np.savez_compressed(d_tr_cache, embeddings=train_bb)\n    np.savez_compressed(d_va_cache, embeddings=val_bb)\n    del backbone; torch.cuda.empty_cache()\n    print(f\"Extracted and cached: train={train_bb.shape}  val={val_bb.shape}\")\n\nprint(f\"Backbone output dim: {dinov2_dim}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:10:08.620120Z","iopub.execute_input":"2026-03-12T18:10:08.620404Z","iopub.status.idle":"2026-03-12T18:32:46.020928Z","shell.execute_reply.started":"2026-03-12T18:10:08.620373Z","shell.execute_reply":"2026-03-12T18:32:46.020185Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Architecture and Training Utilities","metadata":{}},{"cell_type":"code","source":"class EmbeddingProjection(nn.Module):\n    \"\"\"Projection head — width/depth/dropout are configurable per stage.\"\"\"    \n    def __init__(self, input_dim, hidden_dim=512, output_dim=256, dropout=0.3):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(input_dim, hidden_dim),\n            nn.BatchNorm1d(hidden_dim),\n            nn.ReLU(inplace=True),\n            nn.Dropout(dropout),\n            nn.Linear(hidden_dim, output_dim),\n            nn.BatchNorm1d(output_dim),\n        )\n        for m in self.modules():\n            if isinstance(m, nn.Linear):\n                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n                if m.bias is not None: nn.init.constant_(m.bias, 0)\n            elif isinstance(m, nn.BatchNorm1d):\n                nn.init.constant_(m.weight, 1); nn.init.constant_(m.bias, 0)\n\n    def forward(self, x):         return self.net(x)\n    def get_embeddings(self, x):  return F.normalize(self.net(x), p=2, dim=1)\n\n\nclass ArcFaceHead(nn.Module):\n    \"\"\"Additive Angular Margin — margin and scale are configurable.\"\"\"   \n    def __init__(self, emb_dim, num_classes, margin=0.5, scale=64.0):\n        super().__init__()\n        self.scale = scale\n        self.cos_m  = math.cos(margin); self.sin_m = math.sin(margin)\n        self.th     = math.cos(math.pi - margin)\n        self.mm     = math.sin(math.pi - margin) * margin\n        self.W      = nn.Parameter(torch.FloatTensor(num_classes, emb_dim))\n        nn.init.xavier_uniform_(self.W)\n\n    def forward(self, emb, labels):\n        emb_n = F.normalize(emb, p=2, dim=1)\n        w_n   = F.normalize(self.W, p=2, dim=1)\n        cos   = torch.clamp(F.linear(emb_n, w_n), -1, 1)\n        sin   = (1 - cos**2).clamp(min=1e-12).sqrt()\n        phi   = cos * self.cos_m - sin * self.sin_m\n        phi   = torch.where(cos > self.th, phi, cos - self.mm)\n        oh    = torch.zeros_like(cos).scatter_(1, labels.view(-1,1).long(), 1)\n        return (oh * phi + (1-oh) * cos) * self.scale\n\n\nclass EmbeddingDataset(torch.utils.data.Dataset):\n    def __init__(self, embs, labels):\n        self.embs = torch.FloatTensor(embs)\n        self.labels = torch.LongTensor(labels)\n    def __len__(self):          return len(self.labels)\n    def __getitem__(self, idx): return self.embs[idx], self.labels[idx]\n\n\ndef map_from_embeddings(emb, labels):\n    emb  = emb / (np.linalg.norm(emb, axis=1, keepdims=True) + 1e-12)\n    sim  = emb @ emb.T; np.fill_diagonal(sim, -1)\n    id_aps = defaultdict(list)\n    for q in range(len(labels)):\n        ql  = labels[q]; is_m = (labels==ql).astype(int); is_m[q] = 0\n        n_pos = is_m.sum()\n        if n_pos == 0: continue\n        order = np.argsort(-sim[q]); sm = is_m[order]\n        cum   = np.cumsum(sm); prec = cum / np.arange(1, len(sm)+1)\n        id_aps[ql].append(float(np.sum(prec*sm)/n_pos))\n    return float(np.mean([np.mean(v) for v in id_aps.values()]))\n\n\nprint(\"Architecture defined ✓  (EmbeddingProjection, ArcFaceHead, map_from_embeddings)\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:32:46.022046Z","iopub.execute_input":"2026-03-12T18:32:46.022412Z","iopub.status.idle":"2026-03-12T18:32:46.038261Z","shell.execute_reply.started":"2026-03-12T18:32:46.022385Z","shell.execute_reply":"2026-03-12T18:32:46.037580Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def run_config(tag, margin, scale, emb_dim, hidden_dim, dropout,\n               wandb_prefix, log_curve=False):\n    \"\"\"\n    Train one ArcFace configuration on cached DINOv2 embeddings.\n    Returns best val mAP. All hyperparameters are explicit arguments.\n    \"\"\"\n    set_seed(config[\"seed\"])\n\n    proj  = EmbeddingProjection(dinov2_dim, hidden_dim, emb_dim, dropout).to(device)\n    head  = ArcFaceHead(emb_dim, num_classes, margin, scale).to(device)\n    crit  = nn.CrossEntropyLoss()\n\n    train_loader = DataLoader(\n        EmbeddingDataset(train_bb, train_labels_enc),\n        batch_size=config[\"batch_size\"], shuffle=True, drop_last=True)\n    val_loader = DataLoader(\n        EmbeddingDataset(val_bb, val_labels_enc),\n        batch_size=config[\"batch_size\"], shuffle=False)\n\n    opt   = torch.optim.AdamW(\n        list(proj.parameters()) + list(head.parameters()),\n        lr=config[\"learning_rate\"], weight_decay=config[\"weight_decay\"])\n    sched = torch.optim.lr_scheduler.ReduceLROnPlateau(\n        opt, mode='min', factor=0.5, patience=5)\n\n    best_proj  = {k: v.clone() for k, v in proj.state_dict().items()}\n    best_head  = {k: v.clone() for k, v in head.state_dict().items()}\n    best_vloss = float('inf')\n    best_map   = 0.0\n    best_ep    = 0\n    patience_c = 0\n    val_maps   = []\n\n    for epoch in range(config[\"num_epochs\"]):\n        # Train\n        proj.train(); head.train(); tr_loss = 0.0\n        for emb, labels in train_loader:\n            emb, labels = emb.to(device), labels.to(device)\n            loss = crit(head(proj(emb), labels), labels)\n            opt.zero_grad(); loss.backward(); opt.step()\n            tr_loss += loss.item()\n        tr_loss /= len(train_loader)\n\n        # Validate\n        proj.eval(); head.eval(); va_loss = 0.0\n        with torch.no_grad():\n            for emb, labels in val_loader:\n                emb, labels = emb.to(device), labels.to(device)\n                va_loss += crit(head(proj(emb), labels), labels).item()\n        va_loss /= len(val_loader)\n\n        # mAP\n        with torch.no_grad():\n            ve = proj.get_embeddings(\n                torch.FloatTensor(val_bb).to(device)).cpu().numpy()\n        va_map = map_from_embeddings(ve, val_labels_str)\n        val_maps.append(va_map)\n        sched.step(va_loss)\n\n        if log_curve:\n            wandb.log({f\"{wandb_prefix}/epoch\": epoch+1,\n                       f\"{wandb_prefix}/train_loss\": tr_loss,\n                       f\"{wandb_prefix}/val_loss\":   va_loss,\n                       f\"{wandb_prefix}/val_map\":    va_map})\n\n        loss_ok    = (va_loss == va_loss)\n        map_better = va_map > best_map\n        if loss_ok and (va_loss < best_vloss): best_vloss = va_loss\n        if map_better or (loss_ok and va_loss < best_vloss):\n            if map_better:\n                best_map = va_map; best_ep = epoch+1\n                best_proj = {k: v.clone() for k, v in proj.state_dict().items()}\n                best_head = {k: v.clone() for k, v in head.state_dict().items()}\n            patience_c = 0\n        else:\n            patience_c += 1\n            if patience_c >= config[\"patience\"]: break\n\n    wandb.log({f\"{wandb_prefix}/best_map\":   best_map,\n               f\"{wandb_prefix}/best_epoch\": best_ep,\n               f\"{wandb_prefix}/gain\":       best_map - EXP4_BASELINE_MAP})\n\n    # Restore best and return val embeddings\n    proj.load_state_dict(best_proj)\n    proj.eval()\n    with torch.no_grad():\n        best_val_emb = proj.get_embeddings(\n            torch.FloatTensor(val_bb).to(device)).cpu().numpy()\n\n    return best_map, best_ep, best_val_emb, proj\n\n\nprint(\"run_config() training function defined ✓\")\nprint(\"Each call: trains from scratch, in-memory checkpoint, returns best mAP\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:32:46.039238Z","iopub.execute_input":"2026-03-12T18:32:46.039493Z","iopub.status.idle":"2026-03-12T18:32:46.056071Z","shell.execute_reply.started":"2026-03-12T18:32:46.039464Z","shell.execute_reply":"2026-03-12T18:32:46.055503Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4. Stage 1 — Margin × Scale Grid Sweep\n\nFixing: embedding_dim=256, hidden_dim=512, dropout=0.3  (Exp 4 defaults)\nSweeping: margin ∈ {0.2, 0.3, 0.4, 0.5, 0.6, 0.7} × scale ∈ {48, 64, 80, 128}\n","metadata":{}},{"cell_type":"code","source":"stage1_results = []\ntotal_s1 = len(STAGE1_MARGINS) * len(STAGE1_SCALES)\ndone = 0\n\nprint(f\"Stage 1: {total_s1} configs  (dim=256, hidden=512, dropout=0.3)\")\nprint(f\"Baseline: mAP={EXP4_BASELINE_MAP} (margin=0.5, scale=64)\")\nprint()\n\nfor margin in STAGE1_MARGINS:\n    for scale in STAGE1_SCALES:\n        done += 1\n        tag = f\"m{str(margin).replace('.','')}_s{int(scale)}\"\n        t0  = time.perf_counter()\n        val_map, best_ep, _, _ = run_config(\n            tag=tag, margin=margin, scale=scale,\n            emb_dim=256, hidden_dim=512, dropout=0.3,\n            wandb_prefix=f\"stage1/{tag}\")\n        elapsed = time.perf_counter() - t0\n        gain    = val_map - EXP4_BASELINE_MAP\n\n        stage1_results.append({\n            \"margin\": margin, \"scale\": scale,\n            \"val_map\": val_map, \"gain\": gain,\n            \"best_epoch\": best_ep, \"time_s\": elapsed,\n        })\n        marker = \" ★ BEST\" if val_map == max(r[\"val_map\"] for r in stage1_results) else \"\"\n        print(f\"  [{done:2d}/{total_s1}] margin={margin:.1f} scale={int(scale):3d} | \"\n              f\"mAP={val_map:.4f} ({gain:+.4f}) | ep={best_ep:3d} | \"\n              f\"{elapsed/60:.1f}min{marker}\")\n\nstage1_results.sort(key=lambda r: -r[\"val_map\"])\nbest_s1 = stage1_results[0]\nprint(f\"\\nStage 1 best: margin={best_s1['margin']}, scale={int(best_s1['scale'])} \"\n      f\"→ mAP={best_s1['val_map']:.4f} ({best_s1['gain']:+.4f} vs Exp4)\")\n\nwandb.log({\"stage1/best_margin\":  best_s1[\"margin\"],\n           \"stage1/best_scale\":   best_s1[\"scale\"],\n           \"stage1/best_map\":     best_s1[\"val_map\"],\n           \"stage1/best_gain\":    best_s1[\"gain\"]})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:32:46.056927Z","iopub.execute_input":"2026-03-12T18:32:46.057152Z","iopub.status.idle":"2026-03-12T18:37:12.785895Z","shell.execute_reply.started":"2026-03-12T18:32:46.057132Z","shell.execute_reply":"2026-03-12T18:37:12.785262Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Heatmap: margin × scale\ngrid = np.zeros((len(STAGE1_MARGINS), len(STAGE1_SCALES)))\nfor r in stage1_results:\n    i = STAGE1_MARGINS.index(r[\"margin\"])\n    j = STAGE1_SCALES.index(r[\"scale\"])\n    grid[i, j] = r[\"val_map\"]\n\nfig, axes = plt.subplots(1, 2, figsize=(14, 5))\n\nax = axes[0]\nim = ax.imshow(grid, cmap='YlGn', aspect='auto')\nax.set_yticks(range(len(STAGE1_MARGINS)))\nax.set_yticklabels([f\"m={m}\" for m in STAGE1_MARGINS])\nax.set_xticks(range(len(STAGE1_SCALES)))\nax.set_xticklabels([f\"s={int(s)}\" for s in STAGE1_SCALES])\nax.set_title(\"Stage 1: Val mAP — Margin × Scale\", fontweight='bold')\nax.set_xlabel(\"Scale\"); ax.set_ylabel(\"Margin\")\nfor i in range(len(STAGE1_MARGINS)):\n    for j in range(len(STAGE1_SCALES)):\n        best_m = \"★\" if (STAGE1_MARGINS[i]==best_s1[\"margin\"] and\n                          STAGE1_SCALES[j]==best_s1[\"scale\"]) else \"\"\n        ax.text(j, i, f\"{grid[i,j]:.4f}{best_m}\",\n                ha='center', va='center', fontsize=8,\n                color='white' if grid[i,j] > grid.mean() else 'black')\nplt.colorbar(im, ax=ax, label='Val mAP')\n\n# Gain heatmap (vs Exp4 baseline)\ngain_grid = grid - EXP4_BASELINE_MAP\nax = axes[1]\nvmax = max(abs(gain_grid.min()), abs(gain_grid.max()))\nim2 = ax.imshow(gain_grid, cmap='RdYlGn', aspect='auto', vmin=-vmax, vmax=vmax)\nax.set_yticks(range(len(STAGE1_MARGINS))); ax.set_yticklabels([f\"m={m}\" for m in STAGE1_MARGINS])\nax.set_xticks(range(len(STAGE1_SCALES)));  ax.set_xticklabels([f\"s={int(s)}\" for s in STAGE1_SCALES])\nax.set_title(f\"Stage 1: Gain vs Exp4 baseline ({EXP4_BASELINE_MAP})\", fontweight='bold')\nax.set_xlabel(\"Scale\"); ax.set_ylabel(\"Margin\")\nfor i in range(len(STAGE1_MARGINS)):\n    for j in range(len(STAGE1_SCALES)):\n        ax.text(j, i, f\"{gain_grid[i,j]:+.4f}\",\n                ha='center', va='center', fontsize=8)\nplt.colorbar(im2, ax=ax, label='Gain vs Exp4')\n\nplt.tight_layout()\nwandb.log({\"stage1/heatmap\": wandb.Image(fig)})\nplt.savefig('/kaggle/working/stage1_heatmap.png', dpi=150, bbox_inches='tight')\nplt.show()\n\nwandb.log({\"stage1/results_table\": wandb.Table(\n    dataframe=pd.DataFrame([{k:v for k,v in r.items() if k != 'time_s'}\n                             for r in stage1_results]))})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:37:12.786609Z","iopub.execute_input":"2026-03-12T18:37:12.786888Z","iopub.status.idle":"2026-03-12T18:37:15.203712Z","shell.execute_reply.started":"2026-03-12T18:37:12.786853Z","shell.execute_reply":"2026-03-12T18:37:15.203147Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 5. Stage 2 — Embedding Dimension Sweep\n\nFixing: best margin+scale from Stage 1, hidden_dim=512, dropout=0.3\nSweeping: embedding_dim ∈ {128, 256, 512, 1024}\n\nDINOv2 outputs 1024-dim features. Compressing to 256 (Exp 4 default) may discard\ndiscriminative information. Larger output dimensions trade off retrieval speed\nfor representational capacity.\n","metadata":{}},{"cell_type":"code","source":"stage2_results = []\nbest_margin = best_s1[\"margin\"]\nbest_scale  = best_s1[\"scale\"]\n\nprint(f\"Stage 2: embedding_dim sweep\")\nprint(f\"  Fixed: margin={best_margin}, scale={int(best_scale)} (Stage 1 best)\")\nprint(f\"  Fixed: hidden=512, dropout=0.3\")\nprint()\n\nfor emb_dim in STAGE2_EMB_DIMS:\n    tag = f\"dim{emb_dim}\"\n    t0  = time.perf_counter()\n    val_map, best_ep, _, _ = run_config(\n        tag=tag, margin=best_margin, scale=best_scale,\n        emb_dim=emb_dim, hidden_dim=512, dropout=0.3,\n        wandb_prefix=f\"stage2/{tag}\")\n    elapsed = time.perf_counter() - t0\n    gain    = val_map - EXP4_BASELINE_MAP\n\n    stage2_results.append({\n        \"emb_dim\": emb_dim, \"val_map\": val_map,\n        \"gain\": gain, \"best_epoch\": best_ep,\n    })\n    marker = \" ★ BEST\" if val_map == max(r[\"val_map\"] for r in stage2_results) else \"\"\n    print(f\"  dim={emb_dim:4d} | mAP={val_map:.4f} ({gain:+.4f}) | \"\n          f\"ep={best_ep:3d} | {elapsed/60:.1f}min{marker}\")\n\nstage2_results.sort(key=lambda r: -r[\"val_map\"])\nbest_s2 = stage2_results[0]\nprint(f\"\\nStage 2 best: emb_dim={best_s2['emb_dim']} → mAP={best_s2['val_map']:.4f}\")\n\nwandb.log({\"stage2/best_emb_dim\": best_s2[\"emb_dim\"],\n           \"stage2/best_map\":     best_s2[\"val_map\"],\n           \"stage2/best_gain\":    best_s2[\"gain\"]})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:37:15.204548Z","iopub.execute_input":"2026-03-12T18:37:15.204852Z","iopub.status.idle":"2026-03-12T18:37:59.925541Z","shell.execute_reply.started":"2026-03-12T18:37:15.204828Z","shell.execute_reply":"2026-03-12T18:37:59.924968Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dims  = [r[\"emb_dim\"] for r in sorted(stage2_results, key=lambda x: x[\"emb_dim\"])]\nmaps  = [r[\"val_map\"] for r in sorted(stage2_results, key=lambda x: x[\"emb_dim\"])]\ngains = [r[\"gain\"]    for r in sorted(stage2_results, key=lambda x: x[\"emb_dim\"])]\n\nfig, ax = plt.subplots(figsize=(8, 5))\nax.plot(dims, maps, 'o-', color='royalblue', linewidth=2.5, markersize=9)\n\nfor d, m, g in zip(dims, maps, gains):\n    ax.annotate(f\"{m:.4f}\\n({g:+.4f})\", (d, m),\n                textcoords=\"offset points\", xytext=(0, 10),\n                ha='center', fontsize=9)\n\nax.axhline(y=EXP4_BASELINE_MAP, color='red', linestyle='--',\n           label=f\"Exp4 baseline (dim=256): {EXP4_BASELINE_MAP:.4f}\")\nax.scatter([best_s2[\"emb_dim\"]], [best_s2[\"val_map\"]], s=300, color='gold',\n           zorder=4, marker='*', label=f\"Best: dim={best_s2['emb_dim']}\")\n\nax.set_xlabel(\"Embedding Dimension\")\nax.set_ylabel(\"Val mAP\")\nax.set_title(f\"Stage 2: Embedding Dimension\\n(margin={best_margin}, scale={int(best_scale)})\")\nax.set_xscale('log', base=2)\nax.set_xticks(dims)\nax.get_xaxis().set_major_formatter(plt.ScalarFormatter())\nax.legend()\nax.grid(True, alpha=0.3)\nplt.tight_layout()\nwandb.log({\"stage2/dim_plot\": wandb.Image(fig)})\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:37:59.926239Z","iopub.execute_input":"2026-03-12T18:37:59.926470Z","iopub.status.idle":"2026-03-12T18:38:00.301621Z","shell.execute_reply.started":"2026-03-12T18:37:59.926441Z","shell.execute_reply":"2026-03-12T18:38:00.301081Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 6. Stage 3 — Hidden Dim + Dropout Sweep\n\nFixing: best margin+scale from Stage 1, best emb_dim from Stage 2\nSweeping: hidden_dim ∈ {512, 1024} × dropout ∈ {0.1, 0.3, 0.5}\n\nA wider hidden layer (1024) increases the projection head capacity.\nLower dropout may help when the backbone features are already rich (DINOv2);\nhigher dropout acts as regularisation for the small identity count.\n","metadata":{}},{"cell_type":"code","source":"stage3_results = []\nbest_emb_dim = best_s2[\"emb_dim\"]\n\nprint(f\"Stage 3: hidden_dim × dropout sweep\")\nprint(f\"  Fixed: margin={best_margin}, scale={int(best_scale)}, emb_dim={best_emb_dim}\")\nprint()\n\nfor cfg3 in STAGE3_CONFIGS:\n    hidden = cfg3[\"hidden_dim\"]; drop = cfg3[\"dropout\"]\n    tag = f\"h{hidden}_d{str(drop).replace('.','')}\"\n    t0  = time.perf_counter()\n    val_map, best_ep, best_val_emb, best_proj = run_config(\n        tag=tag, margin=best_margin, scale=best_scale,\n        emb_dim=best_emb_dim, hidden_dim=hidden, dropout=drop,\n        wandb_prefix=f\"stage3/{tag}\",\n        log_curve=(hidden==512 and drop==0.3))   # log full curve for Exp4-equivalent\n    elapsed = time.perf_counter() - t0\n    gain    = val_map - EXP4_BASELINE_MAP\n\n    stage3_results.append({\n        \"hidden_dim\": hidden, \"dropout\": drop,\n        \"val_map\": val_map,   \"gain\": gain,\n        \"best_epoch\": best_ep,\n        \"val_emb\": best_val_emb,   # keep for final analysis\n        \"proj\": best_proj,\n    })\n    is_exp4_default = (hidden == 512 and drop == 0.3)\n    marker = \" ★ BEST\" if val_map == max(r[\"val_map\"] for r in stage3_results) else \"\"\n    default = \" (Exp4 default)\" if is_exp4_default else \"\"\n    print(f\"  hidden={hidden}, dropout={drop} | \"\n          f\"mAP={val_map:.4f} ({gain:+.4f}) | \"\n          f\"ep={best_ep}{default}{marker}\")\n\nstage3_results.sort(key=lambda r: -r[\"val_map\"])\nbest_s3 = stage3_results[0]\nprint(f\"\\nStage 3 best: hidden={best_s3['hidden_dim']}, dropout={best_s3['dropout']} \"\n      f\"→ mAP={best_s3['val_map']:.4f}\")\n\nwandb.log({\"stage3/best_hidden\":  best_s3[\"hidden_dim\"],\n           \"stage3/best_dropout\": best_s3[\"dropout\"],\n           \"stage3/best_map\":     best_s3[\"val_map\"],\n           \"stage3/best_gain\":    best_s3[\"gain\"]})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:38:00.302445Z","iopub.execute_input":"2026-03-12T18:38:00.302738Z","iopub.status.idle":"2026-03-12T18:39:08.771354Z","shell.execute_reply.started":"2026-03-12T18:38:00.302711Z","shell.execute_reply":"2026-03-12T18:39:08.770728Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 7. Overall Results Summary","metadata":{}},{"cell_type":"code","source":"# Collect best from each stage + Exp4 baseline\nall_summary = [\n    {\"config\": \"Exp4 baseline\",\n     \"margin\": 0.5, \"scale\": 64, \"emb_dim\": 256, \"hidden\": 512, \"dropout\": 0.3,\n     \"val_map\": EXP4_BASELINE_MAP, \"gain\": 0.0},\n    {\"config\": f\"Stage1 best (margin={best_s1['margin']}, scale={int(best_s1['scale'])})\",\n     \"margin\": best_s1[\"margin\"], \"scale\": best_s1[\"scale\"],\n     \"emb_dim\": 256, \"hidden\": 512, \"dropout\": 0.3,\n     \"val_map\": best_s1[\"val_map\"], \"gain\": best_s1[\"gain\"]},\n    {\"config\": f\"Stage2 best (dim={best_s2['emb_dim']})\",\n     \"margin\": best_margin, \"scale\": best_scale,\n     \"emb_dim\": best_s2[\"emb_dim\"], \"hidden\": 512, \"dropout\": 0.3,\n     \"val_map\": best_s2[\"val_map\"], \"gain\": best_s2[\"gain\"]},\n    {\"config\": f\"Stage3 best (hidden={best_s3['hidden_dim']}, drop={best_s3['dropout']})\",\n     \"margin\": best_margin, \"scale\": best_scale,\n     \"emb_dim\": best_emb_dim,\n     \"hidden\": best_s3[\"hidden_dim\"], \"dropout\": best_s3[\"dropout\"],\n     \"val_map\": best_s3[\"val_map\"], \"gain\": best_s3[\"gain\"]},\n]\nbest_overall = max(all_summary, key=lambda r: r[\"val_map\"])\n\nprint(\"=\" * 70)\nprint(f\"{'Config':45s} | {'mAP':>6} | {'Gain':>7}\")\nprint(\"-\" * 70)\nfor r in all_summary:\n    marker = \"  << OVERALL BEST\" if r[\"config\"] == best_overall[\"config\"] else \"\"\n    print(f\"  {r['config']:43s} | {r['val_map']:.4f} | {r['gain']:+.4f}{marker}\")\nprint(\"=\" * 70)\n\nprint(f\"\\nBest configuration:\")\nprint(f\"  margin={best_overall['margin']}, scale={int(best_overall['scale'])}, \"\n      f\"emb_dim={best_overall['emb_dim']}, \"\n      f\"hidden={best_overall['hidden']}, dropout={best_overall['dropout']}\")\nprint(f\"  Val mAP: {best_overall['val_map']:.4f}  (Exp4 baseline: {EXP4_BASELINE_MAP:.4f})\")\nprint(f\"  Net gain: {best_overall['gain']:+.4f}\")\n\nwandb.log({\n    \"results/best_overall_map\":  best_overall[\"val_map\"],\n    \"results/best_overall_gain\": best_overall[\"gain\"],\n    \"results/best_margin\":       best_overall[\"margin\"],\n    \"results/best_scale\":        best_overall[\"scale\"],\n    \"results/best_emb_dim\":      best_overall[\"emb_dim\"],\n    \"results/best_hidden\":       best_overall[\"hidden\"],\n    \"results/best_dropout\":      best_overall[\"dropout\"],\n    \"results/summary_table\":     wandb.Table(dataframe=pd.DataFrame(all_summary)),\n})\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:39:08.773621Z","iopub.execute_input":"2026-03-12T18:39:08.774494Z","iopub.status.idle":"2026-03-12T18:39:08.935359Z","shell.execute_reply.started":"2026-03-12T18:39:08.774461Z","shell.execute_reply":"2026-03-12T18:39:08.934690Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Final comparison bar chart\nconfigs = [r[\"config\"].replace(\" (\", \"\\n(\") for r in all_summary]   # \\\\n → \\n\nmaps    = [r[\"val_map\"] for r in all_summary]\ncolors  = ['lightgray' if r[\"config\"] == \"Exp4 baseline\"\n           else ('gold' if r[\"config\"] == best_overall[\"config\"]\n           else 'steelblue') for r in all_summary]\n\nfig, ax = plt.subplots(figsize=(12, 5))\nbars = ax.bar(configs, maps, color=colors, edgecolor='black', linewidth=0.8)\n\nfor bar, val, r in zip(bars, maps, all_summary):\n    ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.001,\n            f\"{val:.4f}\\n({r['gain']:+.4f})\",              # \\\\n → \\n\n            ha='center', va='bottom', fontsize=9, fontweight='bold')\n\nax.axhline(y=EXP4_BASELINE_MAP, color='red', linestyle='--', linewidth=1.5,\n           label=f\"Exp4 baseline: {EXP4_BASELINE_MAP:.4f}\")\nax.set_ylabel(\"Val mAP (Identity-Balanced)\")\nax.set_title(\"Experiment 13: ArcFace Sweep — Best per Stage\\n\"   # raw newline → \\n\n             \"(gold = overall best | gain vs Exp4 shown in brackets)\")\nax.legend()\nax.grid(True, alpha=0.3, axis='y')\nax.set_ylim(min(maps) * 0.97, max(maps) * 1.05)\nplt.tight_layout()\nwandb.log({\"results/summary_bar\": wandb.Image(fig)})\nplt.savefig('/kaggle/working/exp13_summary.png', dpi=150, bbox_inches='tight')\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:39:08.936138Z","iopub.execute_input":"2026-03-12T18:39:08.936344Z","iopub.status.idle":"2026-03-12T18:39:09.435166Z","shell.execute_reply.started":"2026-03-12T18:39:08.936323Z","shell.execute_reply":"2026-03-12T18:39:09.434453Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 8. Generate Submission with Best Configuration","metadata":{}},{"cell_type":"code","source":"test_pairs_df = pd.read_csv(config[\"data_dir\"] / \"test.csv\")\ntest_images   = sorted(set(test_pairs_df['query_image'].unique()) |\n                        set(test_pairs_df['gallery_image'].unique()))\ntest_paths    = [config[\"data_dir\"] / \"test/test\" / fn for fn in test_images]\nprint(f\"Test images: {len(test_images)}  |  Test pairs: {len(test_pairs_df)}\")\n\n# Load or extract test backbone embeddings\nte_cache = config[\"cache_dir\"] / \"dinov2_test_embeddings.npz\"\nif te_cache.exists():\n    test_bb = np.load(te_cache)[\"embeddings\"]\n    print(f\"Loaded test backbone cache: {test_bb.shape}\")\nelse:\n    print(\"Extracting DINOv2 test embeddings...\")\n    backbone = timm.create_model(config[\"dinov2_model\"], pretrained=True,\n                                  num_classes=0, img_size=518)\n    backbone.eval()\n    for p in backbone.parameters(): p.requires_grad = False\n    backbone.to(device)\n    preprocess = transforms.Compose([\n        transforms.Resize((518,518)), transforms.ToTensor(),\n        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225]),\n    ])\n    out = []\n    for i in tqdm(range(0, len(test_paths), 32), desc=\"DINOv2 test\"):\n        tensors = []\n        for p in test_paths[i:i+32]:\n            try:    tensors.append(preprocess(Image.open(p).convert(\"RGB\")))\n            except: tensors.append(torch.zeros(3,518,518))\n        out.append(backbone(torch.stack(tensors).to(device)).cpu().numpy())\n    test_bb = np.vstack(out)\n    np.savez_compressed(te_cache, embeddings=test_bb)\n    del backbone; torch.cuda.empty_cache()\n    print(f\"Extracted: {test_bb.shape}\")\n\n# Re-train best overall config on full data to get test projection\nprint(f\"\\nApplying best config: {best_overall['config']}\")\n_, _, _, best_proj_final = run_config(\n    tag=\"final\",\n    margin=best_overall[\"margin\"],\n    scale=float(best_overall[\"scale\"]),\n    emb_dim=best_overall[\"emb_dim\"],\n    hidden_dim=best_overall[\"hidden\"],\n    dropout=best_overall[\"dropout\"],\n    wandb_prefix=\"final_retrain\",\n)\n\nwith torch.no_grad():\n    test_emb = best_proj_final.get_embeddings(\n        torch.FloatTensor(test_bb).to(device)).cpu().numpy()\ntest_emb = test_emb / (np.linalg.norm(test_emb, axis=1, keepdims=True) + 1e-12)\n\nimg_to_idx = {img: i for i, img in enumerate(test_images)}\nsims = np.clip([\n    float(np.dot(test_emb[img_to_idx[r['query_image']]],\n                 test_emb[img_to_idx[r['gallery_image']]]))\n    for _, r in tqdm(test_pairs_df.iterrows(),\n                      total=len(test_pairs_df), desc=\"Submission\")\n], 0.0, 1.0)\n\npd.DataFrame({'row_id': test_pairs_df['row_id'], 'similarity': sims})  .to_csv('/kaggle/working/submission.csv', index=False)\nprint(f\"Submission saved.  sim: min={sims.min():.4f} max={sims.max():.4f} mean={sims.mean():.4f}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:39:09.436164Z","iopub.execute_input":"2026-03-12T18:39:09.436524Z","iopub.status.idle":"2026-03-12T18:43:51.604961Z","shell.execute_reply.started":"2026-03-12T18:39:09.436500Z","shell.execute_reply":"2026-03-12T18:43:51.604375Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sub_art = wandb.Artifact(\"submission\", type=\"submission\",\n    description=f\"ArcFace sweep best: {best_overall['config']}\")\nsub_art.add_file(\"/kaggle/working/submission.csv\")\nwandb.log_artifact(sub_art)\nwandb.finish()\nprint(\"W&B run completed.\")\nprint(f\"\\nFinal result:  {best_overall['val_map']:.4f}  \"\n      f\"({'improved' if best_overall['gain'] > 0 else 'no improvement'} \"\n      f\"vs Exp4 baseline of {EXP4_BASELINE_MAP:.4f}, gain={best_overall['gain']:+.4f})\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-12T18:43:51.605858Z","iopub.execute_input":"2026-03-12T18:43:51.606102Z","iopub.status.idle":"2026-03-12T18:43:52.679104Z","shell.execute_reply.started":"2026-03-12T18:43:51.606071Z","shell.execute_reply":"2026-03-12T18:43:52.678484Z"}},"outputs":[],"execution_count":null}]}