File size: 5,715 Bytes
a3454ce | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | """
Evaluate OMNI encoders on the THUNDER benchmark (https://github.com/MICS-Lab/thunder).
Setup (once):
pip install thunder-bench
export THUNDER_BASE_DATA_FOLDER=/path/to/thunder_data # empty folder, THUNDER fills it
The requested datasets are downloaded and their splits generated automatically
(skip this with --no-download once they are in place).
Usage (Python API):
python eval_thunder.py --model base --datasets bach mhist --tasks knn linear_probing
python eval_thunder.py --model large --datasets pannuke --tasks segmentation
python eval_thunder.py --model sofieneb/omni-small --datasets bach --tasks knn simple_shot
Usage (THUNDER CLI, which instantiates the class below with its defaults):
OMNI_MODEL=base thunder benchmark custom:eval_thunder.py bach knn
OMNI_MODEL=base thunder benchmark custom:eval_thunder.py bach linear_probing --loading-mode=embedding_pre_loading
Results are written to:
$THUNDER_BASE_DATA_FOLDER/outputs/res/<dataset>/<model>/<task>/<adaptation_type>/outputs.json
and can be summarised with:
thunder results-summary
"""
import argparse
import os
import torch
from thunder import benchmark, download_datasets, generate_splits
from thunder.models import PretrainedModel
from transformers import AutoModel
# Short names -> Hugging Face repos
OMNI_MODELS = {
"tiny": "sofieneb/omni-tiny",
"small": "sofieneb/omni-small",
"base": "sofieneb/omni-base",
"large": "sofieneb/omni-large",
}
DEFAULT_MODEL = "tiny"
# Tasks that train on pre-computed embeddings (see the THUNDER examples)
PRELOADED_TASKS = {"linear_probing", "segmentation", "simple_shot"}
def resolve_model(model: str) -> str:
"""Map 'tiny' / 'small' / 'base' / 'large' to a repo id; pass anything else through (repo id or local folder)."""
return OMNI_MODELS.get(model, model)
class OmniThunder(PretrainedModel):
"""
THUNDER wrapper around an OMNI encoder.
Linear probing / kNN / SimpleShot use the average of the teacher tokens;
segmentation uses the patch tokens.
Args:
model (str): 'tiny', 'small', 'base', 'large', a Hugging Face repo id or a local folder.
Defaults to the OMNI_MODEL environment variable, then to 'base'.
"""
def __init__(self, model: str | None = None):
super().__init__()
repo_id = resolve_model(model or os.environ.get("OMNI_MODEL", DEFAULT_MODEL))
self.model = AutoModel.from_pretrained(repo_id, trust_remote_code=True).eval()
self.name = repo_id.rstrip("/").split("/")[-1] # used by THUNDER to name output folders
self.emb_dim = self.model.config.embed_dim
self.vlm = False
self._transform = self.model.get_transform()
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x (torch.Tensor): Images of shape (B, 3, 224, 224).
Returns:
torch.Tensor: Tile embeddings of shape (B, D).
"""
return self.model(x)
def get_transform(self):
"""Returns the image preprocessing used by OMNI."""
return self._transform
def get_linear_probing_embeddings(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x (torch.Tensor): Images of shape (B, 3, 224, 224).
Returns:
torch.Tensor: Average of the teacher tokens, shape (B, D).
"""
return self.model(x)
def get_segmentation_embeddings(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x (torch.Tensor): Images of shape (B, 3, 224, 224).
Returns:
torch.Tensor: Patch tokens of shape (B, 256, D).
"""
return self.model(x, return_patch_tokens=True)["patch_tokens"]
def main():
parser = argparse.ArgumentParser(description="Evaluate OMNI encoders on THUNDER.")
parser.add_argument(
"--model",
default=os.environ.get("OMNI_MODEL", DEFAULT_MODEL),
help="tiny, small, base, large, a Hugging Face repo id or a local folder (default: base).",
)
parser.add_argument("--datasets", nargs="+", default=["bach"], help="THUNDER dataset names.")
parser.add_argument(
"--tasks",
nargs="+",
default=["knn", "linear_probing"],
help="THUNDER tasks: knn, linear_probing, simple_shot, segmentation, "
"transformation_invariance, adversarial_attack.",
)
parser.add_argument(
"--no-download",
dest="download",
action="store_false",
help="Skip downloading the datasets and generating their splits (done by default).",
)
parser.add_argument(
"--no-recompute",
action="store_true",
help="Re-use embeddings already computed by a previous run.",
)
args = parser.parse_args()
if not os.environ.get("THUNDER_BASE_DATA_FOLDER"):
raise SystemExit("Set THUNDER_BASE_DATA_FOLDER first, e.g. export THUNDER_BASE_DATA_FOLDER=/path/to/thunder_data")
if args.download:
download_datasets(args.datasets)
generate_splits(args.datasets)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = OmniThunder(args.model).to(device).eval()
print(f"Model: {model.name} (dim {model.emb_dim}) on {device}")
for dataset in args.datasets:
# Embeddings are computed once, then re-used by every task.
benchmark(model, dataset=dataset, task="pre_computing_embeddings", recomp_embs=not args.no_recompute)
for task in args.tasks:
extra = {"loading_mode": "embedding_pre_loading"} if task in PRELOADED_TASKS else {}
benchmark(model, dataset=dataset, task=task, retrain_model=True, **extra)
if __name__ == "__main__":
main() |