Upload ensemble_eval.py with huggingface_hub

ensemble_eval.py

@@ -0,0 +1,159 @@
+"""Multi-model ensemble + proper TTA — averages softmax across 1+ checkpoints.
+
+Usage:
+  python ensemble_eval.py best_ema.pt
+  python ensemble_eval.py best_ema_v2.pt convnextv2_best_ema.pt
+  python ensemble_eval.py best_ema_v2.pt swa_v2.pt convnextv2_best_ema.pt convnextv2_swa.pt
+
+Each checkpoint is loaded with its own model_name + img_size from its dict
+(so EVA-02 @ 448 and ConvNeXt V2 @ 384 mix freely). Each runs TTA = identity +
+hflip + vflip + 2 scale crops, then per-image softmaxes are averaged.
+
+Saves predictions.csv and prints classification_report.
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+import numpy as np
+import timm
+import torch
+from sklearn.metrics import classification_report, f1_score
+from timm.data import create_transform
+from torch.utils.data import DataLoader
+from torchvision import datasets
+
+BASE_EXTRACT_DIR = "./dermnet-skin40-cleaned-dataset"
+DATA_DIR = os.path.join(BASE_EXTRACT_DIR, "kaggle/working/Merged_Dermnet_Skin40")
+TEST_DIR = os.path.join(DATA_DIR, "test")
+BATCH_SIZE = 16
+NUM_WORKERS = 4
+MEAN, STD = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
+
+os.environ["HIP_VISIBLE_DEVICES"] = "0"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def make_loader(img_size, crop_pct):
+    """Each (img_size, crop_pct) combo gets its own loader for proper TTA."""
+    tf = create_transform(input_size=img_size, is_training=False,
+        crop_pct=crop_pct, interpolation="bicubic", mean=MEAN, std=STD)
+    ds = datasets.ImageFolder(TEST_DIR, transform=tf)
+    return ds, DataLoader(ds, batch_size=BATCH_SIZE, shuffle=False,
+        num_workers=NUM_WORKERS, pin_memory=True)
+
+
+def load_model(ckpt_path):
+    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    model_name = ckpt["model_name"]
+    img_size = ckpt["img_size"]
+    sd = ckpt["model_state_dict"]
+    # Infer num_classes from head weight
+    head_keys = [k for k in sd.keys() if k.endswith("head.weight") or k.endswith("fc.weight")]
+    num_classes = sd[head_keys[0]].shape[0] if head_keys else 23
+    model = timm.create_model(model_name, pretrained=False,
+        num_classes=num_classes, img_size=img_size if "vit" in model_name or "eva" in model_name else None)
+    if "img_size" in timm.create_model.__code__.co_varnames:
+        pass
+    # Some non-ViT models reject img_size kwarg — handle gracefully
+    try:
+        model = timm.create_model(model_name, pretrained=False,
+            num_classes=num_classes, img_size=img_size)
+    except TypeError:
+        model = timm.create_model(model_name, pretrained=False, num_classes=num_classes)
+    model.load_state_dict(sd)
+    model = model.to(device, memory_format=torch.channels_last).eval()
+    print(f"  loaded {ckpt_path}  ({model_name} @ {img_size}, "
+          f"prev acc={ckpt.get('val_acc', 0)*100:.2f}%)")
+    return model, img_size
+
+
+@torch.no_grad()
+def tta_softmax(model, img_size):
+    """4-augmentation TTA: identity + hflip @ crop_pct=0.95, plus same pair @ crop_pct=1.0."""
+    aggregated = None
+    targets = None
+    classes = None
+
+    for crop_pct in [0.95, 1.0]:
+        ds, loader = make_loader(img_size, crop_pct)
+        if classes is None:
+            classes = ds.classes
+
+        batch_softmax = []
+        batch_targets = []
+        for x, y in loader:
+            x = x.to(device, non_blocking=True).to(memory_format=torch.channels_last)
+            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+                # identity
+                p = torch.softmax(model(x), dim=-1).float()
+                # hflip
+                p = p + torch.softmax(model(torch.flip(x, dims=[-1])), dim=-1).float()
+            batch_softmax.append(p.cpu())
+            batch_targets.append(y)
+
+        crop_softmax = torch.cat(batch_softmax)  # 2 augs already summed
+        if aggregated is None:
+            aggregated = crop_softmax
+            targets = torch.cat(batch_targets)
+        else:
+            aggregated = aggregated + crop_softmax
+
+    # Total: 2 crops × 2 flips = 4 augmentations summed
+    aggregated = aggregated / 4.0
+    return aggregated, targets, classes
+
+
+def main():
+    if len(sys.argv) < 2:
+        print("Usage: python ensemble_eval.py <ckpt1> [ckpt2 ...]")
+        sys.exit(1)
+
+    ckpts = sys.argv[1:]
+    print(f"Ensembling {len(ckpts)} checkpoint(s) with TTA (4 augs/model):")
+
+    all_probs = []
+    targets = None
+    classes = None
+
+    for ckpt in ckpts:
+        if not os.path.exists(ckpt):
+            print(f"  SKIP missing: {ckpt}"); continue
+        model, img_size = load_model(ckpt)
+        probs, t, cls = tta_softmax(model, img_size)
+        all_probs.append(probs)
+        if targets is None:
+            targets, classes = t, cls
+        # Free GPU memory between models
+        del model; torch.cuda.empty_cache()
+
+    if not all_probs:
+        print("No valid checkpoints loaded."); sys.exit(1)
+
+    # Equal-weight average across models
+    ensemble = torch.stack(all_probs).mean(0)
+    preds = ensemble.argmax(1)
+
+    acc = (preds == targets).float().mean().item()
+    f1 = f1_score(targets.numpy(), preds.numpy(), average="macro")
+
+    print(f"\n{'='*60}")
+    print(f"Ensemble of {len(all_probs)} model(s) + 4-aug TTA")
+    print(f"{'='*60}")
+    print(f"Accuracy:   {acc*100:.2f}%")
+    print(f"Macro F1:   {f1:.4f}")
+    print(f"\nPer-class report:")
+    print(classification_report(targets.numpy(), preds.numpy(),
+        target_names=classes, digits=3, zero_division=0))
+
+    # Save predictions for further analysis
+    np.savetxt("predictions.csv",
+        np.column_stack([targets.numpy(), preds.numpy(), ensemble.numpy()]),
+        delimiter=",", fmt="%g",
+        header="true,pred," + ",".join(f"p_{c}" for c in classes), comments="")
+    print(f"Saved predictions.csv  ({ensemble.shape[0]} rows)")
+
+
+if __name__ == "__main__":
+    main()