{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":126777,"databundleVersionId":15314950,"isSourceIdPinned":false}],"dockerImageVersionId":31287,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"\"\"\"\nv04: Backbone Exploration\n==========================\nCompare three backbones using the same ArcFace setup as v03.\nGoal: find which architecture captures jaguar spot patterns best.\n\nBackbones tested:\n  efficientnetv2_s      — Fused-MBConv, ~22M params, trains faster than B0\n  convnext_small        — 7x7 depthwise conv, large texture receptive field\n  vit_small_patch16_224 — global self-attention, holistic pattern matching\n\nPerformance optimisations vs v03:\n  1. Image pre-cache: all 1895 images loaded to RAM once (PNG decode happens\n     once, not once per epoch per worker). Eliminates the disk I/O bottleneck\n     that made v02/v03 take 8-10h for ~1895 small images.\n  2. Mixed precision (AMP): torch.cuda.amp halves memory/compute on T4/P100.\n  Together these should cut wall-clock time to ~1.5-2h per backbone.\n\nAll backbones:\n  ArcFace(m=0.5, s=30), embedding_dim=512, img_size=224\n  Stage 1: 5 epochs frozen | Stage 2: 25 epochs full fine-tune\n  Best backbone selected by val mAP, retrained on full data for submission.\n\"\"\"\n\nimport sys, os, time\n\nREPO_URL = 'https://github.com/Smooth-Cactus0/jaguar-re-identification.git'\nREPO_DIR = '/kaggle/working/repo'\n\nif os.path.exists(REPO_DIR):\n    os.system(f'git -C {REPO_DIR} pull --ff-only')\nelse:\n    os.system(f'git clone {REPO_URL} {REPO_DIR}')\nsys.path.insert(0, REPO_DIR)\n\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import OneCycleLR\nfrom torch.cuda.amp import GradScaler, autocast\nfrom pathlib import Path\nimport warnings\nwarnings.filterwarnings('ignore')\n\nfrom src.data      import (JaguarDataset, get_transforms, encode_labels,\n                            get_identity_splits, build_image_cache)\nfrom src.models    import EmbeddingModel\nfrom src.losses    import ArcFaceLoss\nfrom src.inference import extract_embeddings, make_submission\nfrom src.evaluate  import compute_map, print_results, save_benchmark\n\n# -- Paths ---------------------------------------------------------------------\nKAGGLE_INPUT = Path('/kaggle/input/competitions/jaguar-re-id')\nTRAIN_DIR    = KAGGLE_INPUT / 'train' / 'train'\nTEST_DIR     = KAGGLE_INPUT / 'test'  / 'test'\nOUT_DIR      = Path('/kaggle/working/output')\nOUT_DIR.mkdir(exist_ok=True)\n\n# -- Shared config -------------------------------------------------------------\nEMBEDDING_DIM   = 512\nIMG_SIZE        = 224\nBATCH_SIZE      = 64   # AMP enables larger batches; 64 fits all backbones\nN_CLASSES       = 31\nLR_HEAD         = 1e-3\nLR_BACKBONE     = 1e-5\nEPOCHS_S1       = 5\nEPOCHS_S2       = 25\nWEIGHT_DECAY    = 1e-4\nARC_MARGIN      = 0.5\nARC_SCALE       = 30.0\nGRAD_CLIP_NORM  = 1.0\nDEVICE          = 'cuda' if torch.cuda.is_available() else 'cpu'\nUSE_AMP         = DEVICE == 'cuda'   # mixed precision only on GPU\nVERSION         = 'v04'\nCV_FOLD         = 0\n\nBACKBONES = [\n    'tf_efficientnetv2_s',    # TF-converted weights, reliably available in all timm versions\n    'convnext_small',\n    'vit_small_patch16_224',\n]\n\nprint('\\n' + '='*60)\nprint(f'  {VERSION}: Backbone Exploration')\nprint(f'  Comparing : {\", \".join(BACKBONES)}')\nprint(f'  ArcFace   : m={ARC_MARGIN}, s={ARC_SCALE}')\nprint(f'  Device    : {DEVICE}  |  AMP: {USE_AMP}')\nprint(f'  Epochs    : S1={EPOCHS_S1} + S2={EPOCHS_S2} per backbone')\nprint(f'  Batch size: {BATCH_SIZE}')\nprint('='*60)\n\n# -- Load data -----------------------------------------------------------------\nprint('\\n[1] Loading data...')\ntrain_df = pd.read_csv(KAGGLE_INPUT / 'train.csv')\ntest_df  = pd.read_csv(KAGGLE_INPUT / 'test.csv')\nlabels, label_to_idx, idx_to_label = encode_labels(train_df['ground_truth'])\n\nfor fold, train_idx, val_idx in get_identity_splits(train_df, n_splits=5):\n    if fold == CV_FOLD:\n        break\nprint(f'  Train fold: {len(train_idx)} | Val fold: {len(val_idx)}')\n\n# -- Pre-cache ALL images to RAM (eliminates PNG decode bottleneck) ------------\nprint('\\n[2] Pre-caching images to RAM...')\nt_cache = time.time()\nall_train_files  = train_df['filename'].tolist()\ntest_files_flat  = sorted(pd.unique(test_df[['query_image', 'gallery_image']].values.ravel()))\ntrain_cache = build_image_cache(all_train_files, TRAIN_DIR, img_size=IMG_SIZE, verbose=True)\ntest_cache  = build_image_cache(test_files_flat,  TEST_DIR,  img_size=IMG_SIZE, verbose=True)\nprint(f'  Cache built in {time.time()-t_cache:.0f}s')\n\n# Transforms: cached=True skips Resize (already done in build_image_cache)\ntf_train = get_transforms(mode='train', img_size=IMG_SIZE, cached=True)\ntf_val   = get_transforms(mode='val',   img_size=IMG_SIZE, cached=True)\n\n# Build datasets (all share the same cache — zero extra memory overhead)\ntrain_ds = JaguarDataset(\n    filenames=train_df['filename'].iloc[train_idx].tolist(),\n    img_dir=TRAIN_DIR, labels=labels[train_idx], transform=tf_train,\n    image_cache=train_cache)\nval_ds = JaguarDataset(\n    filenames=train_df['filename'].iloc[val_idx].tolist(),\n    img_dir=TRAIN_DIR, labels=labels[val_idx], transform=tf_val,\n    image_cache=train_cache)\nfull_ds = JaguarDataset(\n    filenames=all_train_files, img_dir=TRAIN_DIR, labels=labels,\n    transform=tf_val, image_cache=train_cache)\nfull_train_ds = JaguarDataset(\n    filenames=all_train_files, img_dir=TRAIN_DIR, labels=labels,\n    transform=tf_train, image_cache=train_cache)\ntest_ds_all = JaguarDataset(\n    filenames=test_files_flat, img_dir=TEST_DIR, labels=None,\n    transform=tf_val, image_cache=test_cache)\n\ntrain_loader     = DataLoader(train_ds,     batch_size=BATCH_SIZE, shuffle=True,\n                               num_workers=2, pin_memory=True)\nval_loader       = DataLoader(val_ds,       batch_size=BATCH_SIZE, shuffle=False,\n                               num_workers=2, pin_memory=True)\n\n\n# -- Model + training helpers --------------------------------------------------\n\nclass ArcFaceModel(nn.Module):\n    def __init__(self, backbone_name, embedding_dim, pretrained=True):\n        super().__init__()\n        self.backbone = EmbeddingModel(backbone_name, pretrained=pretrained)\n        native_dim    = self.backbone.out_dim\n        self.proj     = nn.Sequential(\n            nn.Linear(native_dim, embedding_dim, bias=False),\n            nn.BatchNorm1d(embedding_dim),\n        )\n        self.out_dim = embedding_dim\n\n    def forward(self, x):\n        return F.normalize(self.proj(self.backbone(x)), p=2, dim=1)\n\n    def freeze_backbone(self):\n        for p in self.backbone.parameters():\n            p.requires_grad_(False)\n\n    def unfreeze_backbone(self):\n        for p in self.backbone.parameters():\n            p.requires_grad_(True)\n\n\ndef run_epoch(model, arcface_loss, loader, optimizer, scheduler,\n              is_train, device, scaler):\n    \"\"\"Single epoch with AMP (no-op scaler when USE_AMP=False).\"\"\"\n    model.train(is_train)\n    arcface_loss.train(is_train)\n    total_loss, n = 0.0, 0\n    ctx = torch.enable_grad() if is_train else torch.no_grad()\n    with ctx:\n        for imgs, targets in loader:\n            imgs, targets = imgs.to(device), targets.to(device)\n            if is_train:\n                optimizer.zero_grad()\n            with autocast(enabled=USE_AMP):\n                embs = model(imgs)\n                loss = arcface_loss(embs, targets)\n            if is_train:\n                scaler.scale(loss).backward()\n                scaler.unscale_(optimizer)\n                torch.nn.utils.clip_grad_norm_(\n                    list(model.parameters()) + list(arcface_loss.parameters()),\n                    max_norm=GRAD_CLIP_NORM)\n                scaler.step(optimizer)\n                scaler.update()\n                scheduler.step()\n            total_loss += loss.item() * len(imgs)\n            n          += len(imgs)\n    return total_loss / n\n\n\ndef train_backbone(backbone_name):\n    \"\"\"Full two-stage ArcFace training for one backbone. Returns best val mAP.\"\"\"\n    print(f'\\n{\"=\"*55}')\n    print(f'  Backbone: {backbone_name}')\n    print(f'{\"=\"*55}')\n\n    model   = ArcFaceModel(backbone_name, EMBEDDING_DIM).to(DEVICE)\n    arcface = ArcFaceLoss(EMBEDDING_DIM, N_CLASSES, ARC_MARGIN, ARC_SCALE).to(DEVICE)\n    scaler  = GradScaler(enabled=USE_AMP)\n    n_total = sum(p.numel() for p in model.parameters())\n    print(f'  Native dim: {model.backbone.out_dim} | Total params: {n_total:,}')\n\n    # Stage 1: head + ArcFace only\n    model.freeze_backbone()\n    head_params = list(filter(lambda p: p.requires_grad, model.parameters())) + \\\n                  list(arcface.parameters())\n    opt1 = AdamW(head_params, lr=LR_HEAD, weight_decay=WEIGHT_DECAY)\n    sch1 = OneCycleLR(opt1, max_lr=LR_HEAD,\n                      steps_per_epoch=len(train_loader), epochs=EPOCHS_S1)\n    for ep in range(1, EPOCHS_S1 + 1):\n        t0 = time.time()\n        tr = run_epoch(model, arcface, train_loader, opt1, sch1, True,  DEVICE, scaler)\n        va = run_epoch(model, arcface, val_loader,   opt1, sch1, False, DEVICE, scaler)\n        print(f'  S1 ep {ep}/{EPOCHS_S1} | loss {tr:.4f}/{va:.4f} | {time.time()-t0:.0f}s')\n\n    # Stage 2: full fine-tune\n    model.unfreeze_backbone()\n    opt2 = AdamW([\n        {'params': model.backbone.parameters(), 'lr': LR_BACKBONE},\n        {'params': model.proj.parameters(),     'lr': LR_HEAD},\n        {'params': arcface.parameters(),        'lr': LR_HEAD},\n    ], weight_decay=WEIGHT_DECAY)\n    sch2 = OneCycleLR(opt2, max_lr=[LR_BACKBONE, LR_HEAD, LR_HEAD],\n                      steps_per_epoch=len(train_loader), epochs=EPOCHS_S2)\n\n    best_val_map  = 0.0\n    best_state    = None\n\n    for ep in range(1, EPOCHS_S2 + 1):\n        t0 = time.time()\n        tr = run_epoch(model, arcface, train_loader, opt2, sch2, True,  DEVICE, scaler)\n        va = run_epoch(model, arcface, val_loader,   opt2, sch2, False, DEVICE, scaler)\n\n        if ep % 5 == 0 or ep == EPOCHS_S2:\n            val_embs    = extract_embeddings(model, val_ds, BATCH_SIZE, DEVICE,\n                                             l2_normalise=False, desc='')\n            val_results = compute_map(val_embs, labels[val_idx])\n            vm = val_results['map']\n            marker = ''\n            if vm > best_val_map:\n                best_val_map = vm\n                best_state   = {k: v.cpu().clone() for k, v in model.state_dict().items()}\n                marker = ' **'\n            print(f'  S2 ep {ep}/{EPOCHS_S2} | loss {tr:.4f}/{va:.4f} | '\n                  f'mAP {vm:.4f}{marker} | {time.time()-t0:.0f}s')\n        else:\n            print(f'  S2 ep {ep}/{EPOCHS_S2} | loss {tr:.4f}/{va:.4f} | '\n                  f'{time.time()-t0:.0f}s')\n\n    model.load_state_dict(best_state)\n\n    # Full-train mAP\n    full_embs    = extract_embeddings(model, full_ds, BATCH_SIZE, DEVICE,\n                                      l2_normalise=False, desc='Full train')\n    full_results = compute_map(full_embs, labels)\n    print(f'  {backbone_name}: val_mAP={best_val_map:.4f}, full_mAP={full_results[\"map\"]:.4f}')\n\n    return best_val_map, full_results, model\n\n\n# -- Run backbone search -------------------------------------------------------\nresults_all  = {}   # backbone -> (val_map, full_results, model)\nt_total = time.time()\n\nfor backbone in BACKBONES:\n    val_map, full_results, model = train_backbone(backbone)\n    results_all[backbone] = (val_map, full_results, model)\n    # Free GPU memory before next backbone\n    del model\n    torch.cuda.empty_cache()\n\nprint(f'\\nAll backbones trained in {(time.time()-t_total)/60:.0f} min')\n\n# -- Summary -------------------------------------------------------------------\nprint(f'\\n{\"=\"*60}')\nprint(f'  {VERSION} — Backbone Comparison Summary')\nprint(f'  {\"Backbone\":<30} {\"Val mAP\":>8} {\"Full mAP\":>10}')\nprint(f'  {\"-\"*52}')\nbest_backbone = None\nbest_val = 0.0\nfor bb, (vm, fr, _) in results_all.items():\n    marker = ' <-- best' if vm == max(v[0] for v in results_all.values()) else ''\n    print(f'  {bb:<30} {vm:>8.4f} {fr[\"map\"]:>10.4f}{marker}')\n    if vm > best_val:\n        best_val      = vm\n        best_backbone = bb\nprint(f'{\"=\"*60}')\n\n# Save benchmark rows\nfor bb, (vm, fr, _) in results_all.items():\n    save_benchmark(\n        OUT_DIR / 'benchmarks_v04.csv',\n        {\n            'version':        f'{VERSION}_{bb.split(\"_\")[0]}',\n            'backbone':       bb,\n            'loss':           f'ArcFace (m={ARC_MARGIN}, s={ARC_SCALE})',\n            'embedding_dim':  EMBEDDING_DIM,\n            'img_size':       IMG_SIZE,\n            'augmentation':   'none (v05 adds augmentation)',\n            'cv_map_val':     round(vm, 5),\n            'cv_map_full':    round(fr['map'], 5),\n            'best_identity':  idx_to_label[max(fr['per_identity_ap'],\n                                               key=fr['per_identity_ap'].get)],\n            'worst_identity': idx_to_label[min(fr['per_identity_ap'],\n                                               key=fr['per_identity_ap'].get)],\n            'epochs_s1':      EPOCHS_S1,\n            'epochs_s2':      EPOCHS_S2,\n            'notes':          f'Backbone search; best={best_backbone}; AMP+cache',\n        }\n    )\n\n\n# -- Visualisations ------------------------------------------------------------\nprint('\\n[Vis] Generating backbone comparison chart...')\n\nbbs   = list(results_all.keys())\nvmaps = [results_all[b][0]       for b in bbs]\nfmaps = [results_all[b][1]['map'] for b in bbs]\nx     = np.arange(len(bbs))\nw     = 0.35\n\nfig, axes = plt.subplots(1, 2, figsize=(16, 5))\n\naxes[0].bar(x - w/2, vmaps, w, color='lightsteelblue', label='Val fold mAP')\naxes[0].bar(x + w/2, fmaps, w, color='steelblue',      label='Full train mAP')\naxes[0].axhline(0.68287, color='coral', linestyle='--', linewidth=1.5,\n                label='v03 EfficientNet-B0 (full mAP=0.683)')\naxes[0].set_xticks(x)\naxes[0].set_xticklabels([b.replace('_', '\\n') for b in bbs], fontsize=9)\naxes[0].set_ylabel('Identity-balanced mAP'); axes[0].set_ylim(0, 1)\naxes[0].set_title('v04: Backbone Comparison', fontsize=12, fontweight='bold')\naxes[0].legend()\n\n# Per-identity AP for the best backbone\nbest_per_id = {idx_to_label[k]: v\n               for k, v in results_all[best_backbone][1]['per_identity_ap'].items()}\nids_sorted  = sorted(best_per_id, key=lambda x: -best_per_id[x])\nxp = np.arange(len(ids_sorted))\naxes[1].bar(xp, [best_per_id[i] for i in ids_sorted], color='steelblue')\naxes[1].axhline(results_all[best_backbone][1]['map'], color='navy',\n                linestyle='--', linewidth=1,\n                label=f'mAP={results_all[best_backbone][1][\"map\"]:.4f}')\naxes[1].set_xticks(xp)\naxes[1].set_xticklabels(ids_sorted, rotation=45, ha='right', fontsize=7)\naxes[1].set_ylabel('AP'); axes[1].set_ylim(0, 1)\naxes[1].set_title(f'Per-identity AP: {best_backbone}', fontsize=11, fontweight='bold')\naxes[1].legend(fontsize=9)\n\nplt.suptitle('v04: Backbone Search Results', fontsize=13, fontweight='bold')\nplt.tight_layout()\nfig.savefig(OUT_DIR / f'{VERSION}_backbone_comparison.png', dpi=150, bbox_inches='tight')\nplt.close()\nprint('  Saved -> v04_backbone_comparison.png')\n\n\n# -- Full retrain with best backbone + submission ------------------------------\nprint(f'\\n[Final] Retraining best backbone ({best_backbone}) on ALL data...')\n\nfull_loader = DataLoader(full_train_ds, batch_size=BATCH_SIZE, shuffle=True,\n                         num_workers=2, pin_memory=True)\n\nmodel_final   = ArcFaceModel(best_backbone, EMBEDDING_DIM).to(DEVICE)\narcface_final = ArcFaceLoss(EMBEDDING_DIM, N_CLASSES, ARC_MARGIN, ARC_SCALE).to(DEVICE)\nscaler_final  = GradScaler(enabled=USE_AMP)\n\n# Stage 1\nmodel_final.freeze_backbone()\nhead_f = list(filter(lambda p: p.requires_grad, model_final.parameters())) + \\\n         list(arcface_final.parameters())\nopt_f1 = AdamW(head_f, lr=LR_HEAD, weight_decay=WEIGHT_DECAY)\nsch_f1 = OneCycleLR(opt_f1, max_lr=LR_HEAD,\n                    steps_per_epoch=len(full_loader), epochs=EPOCHS_S1)\nfor ep in range(1, EPOCHS_S1 + 1):\n    loss = run_epoch(model_final, arcface_final, full_loader,\n                     opt_f1, sch_f1, True, DEVICE, scaler_final)\n    print(f'  Full S1 ep {ep}/{EPOCHS_S1} | loss {loss:.4f}')\n\n# Stage 2\nmodel_final.unfreeze_backbone()\nopt_f2 = AdamW([\n    {'params': model_final.backbone.parameters(), 'lr': LR_BACKBONE},\n    {'params': model_final.proj.parameters(),     'lr': LR_HEAD},\n    {'params': arcface_final.parameters(),        'lr': LR_HEAD},\n], weight_decay=WEIGHT_DECAY)\nsch_f2 = OneCycleLR(opt_f2, max_lr=[LR_BACKBONE, LR_HEAD, LR_HEAD],\n                    steps_per_epoch=len(full_loader), epochs=EPOCHS_S2)\nfor ep in range(1, EPOCHS_S2 + 1):\n    loss = run_epoch(model_final, arcface_final, full_loader,\n                     opt_f2, sch_f2, True, DEVICE, scaler_final)\n    if ep % 5 == 0:\n        print(f'  Full S2 ep {ep}/{EPOCHS_S2} | loss {loss:.4f}')\n\n# Submission\nmodel_final.eval()\ntest_embs = extract_embeddings(model_final, test_ds_all, BATCH_SIZE, DEVICE,\n                                l2_normalise=False, desc='Test')\nfname_to_idx = {f: i for i, f in enumerate(test_files_flat)}\nmake_submission(\n    test_df=test_df, test_embeddings=test_embs,\n    filename_to_idx=fname_to_idx,\n    output_path=OUT_DIR / f'submission_{VERSION}.csv')\n\ntorch.save(model_final.state_dict(), OUT_DIR / f'model_{VERSION}_{best_backbone}.pth')\n\nprint(f'\\n{\"=\"*60}')\nprint(f'  {VERSION} COMPLETE')\nprint(f'  Best backbone: {best_backbone}  (val_mAP={best_val:.4f})')\nfor bb, (vm, fr, _) in results_all.items():\n    print(f'  {bb:<30}: val={vm:.4f}  full={fr[\"map\"]:.4f}')\nprint(f'  Total time   : {(time.time()-t_total)/60:.0f} min')\nprint(f'  Outputs in   : {OUT_DIR}')\nprint(f'{\"=\"*60}')\nprint('\\nDownload /kaggle/working/output/ zip and push figures + results to GitHub.')","metadata":{"_uuid":"6d035b7d-f589-46f4-a818-8b253d7a08ac","_cell_guid":"b57b848a-cb0d-43c5-8f10-501a63019fb3","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null}]}