Update generalist evaluation

Author: constantinpape

finetuning/generalists/evaluate_generalist.py

@@ -3,31 +3,36 @@
 from subprocess import run
 
 from util import evaluate_checkpoint_for_dataset, ALL_DATASETS, EM_DATASETS, LM_DATASETS
+from micro_sam.evaluation import default_experiment_settings, get_experiment_setting_name
 
 EXPERIMENT_ROOT = "/scratch/projects/nim00007/sam/experiments/generalists"
 CHECKPOINTS = {
+    # Vanilla models
     "vit_b": "/home/nimcpape/.sam_models/sam_vit_b_01ec64.pth",
     "vit_h": "/home/nimcpape/.sam_models/sam_vit_h_4b8939.pth",
+    # Generalist LM models
+    "vit_b_lm": "/scratch-grete/projects/nim00007/sam/models/LM/generalist/v2/vit_b/best.pt",
+    "vit_h_lm": "/scratch-grete/projects/nim00007/sam/models/LM/generalist/v2/vit_h/best.pt",
+    # Generalist EM models
+    "vit_b_em": "/scratch-grete/projects/nim00007/sam/models/EM/generalist/v2/vit_b/best.pt",
+    "vit_h_em": "/scratch-grete/projects/nim00007/sam/models/EM/generalist/v2/vit_h/best.pt",
 }
 
 
-def submit_array_job(model_name, datasets, amg):
+def submit_array_job(model_name, datasets):
     n_datasets = len(datasets)
     cmd = ["sbatch", "-a", f"0-{n_datasets-1}", "evaluate_generalist.sbatch", model_name, "--datasets"]
     cmd.extend(datasets)
-    if amg:
-        cmd.append("--amg")
     run(cmd)
 
 
-def evaluate_dataset_slurm(model_name, dataset, run_amg):
-    max_num_val_images = None
-    if run_amg:
-        if dataset in EM_DATASETS:
-            run_amg = False
-        else:
-            run_amg = True
-            max_num_val_images = 100
+def evaluate_dataset_slurm(model_name, dataset):
+    if dataset in EM_DATASETS:
+        run_amg = False
+        max_num_val_images = None
+    else:
+        run_amg = True
+        max_num_val_images = 64
 
     is_custom_model = model_name not in ("vit_h", "vit_b")
     checkpoint = CHECKPOINTS[model_name]
@@ -52,13 +57,29 @@ def _get_datasets(lm, em):
     return datasets
 
 
+def check_computation(model_name, datasets):
+    prompt_settings = default_experiment_settings()
+    for ds in datasets:
+        experiment_folder = os.path.join(EXPERIMENT_ROOT, model_name, ds)
+        for setting in prompt_settings:
+            setting_name = get_experiment_setting_name(setting)
+            expected_path = os.path.join(experiment_folder, "results", f"{setting_name}.csv")
+            if not os.path.exists(expected_path):
+                print("Missing results for:", expected_path)
+        if ds in LM_DATASETS:
+            expected_path = os.path.join(experiment_folder, "results", "amg.csv")
+            if not os.path.exists(expected_path):
+                print("Missing results for:", expected_path)
+    print("All checks_run")
+
+
 # evaluation on slurm
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("model_name")
+    parser.add_argument("--check", "-c", action="store_true")
     parser.add_argument("--lm", action="store_true")
     parser.add_argument("--em", action="store_true")
-    parser.add_argument("--amg", action="store_true")
     parser.add_argument("--datasets", nargs="+")
     args = parser.parse_args()
 
@@ -67,12 +88,16 @@ def main():
         datasets = _get_datasets(args.lm, args.em)
     assert all(ds in ALL_DATASETS for ds in datasets)
 
+    if args.check:
+        check_computation(args.model_name, datasets)
+        return
+
     job_id = os.environ.get("SLURM_ARRAY_TASK_ID", None)
     if job_id is None:  # this is the main script that submits slurm jobs
-        submit_array_job(args.model_name, datasets, args.amg)
-    else:  # we're in a slurm job and precompute a setting
+        submit_array_job(args.model_name, datasets)
+    else:  # we're in a slurm job
         job_id = int(job_id)
-        evaluate_dataset_slurm(args.model_name, datasets[job_id], args.amg)
+        evaluate_dataset_slurm(args.model_name, datasets[job_id])
 
 
 if __name__ == "__main__":

finetuning/generalists/evaluate_generalist.sbatch

@@ -1,7 +1,7 @@
 #! /bin/bash
 #SBATCH -c 4
 #SBATCH --mem 48G
-#SBATCH -t 720
+#SBATCH -t 2800
 #SBATCH -p grete:shared
 #SBATCH -G A100:1
 #SBATCH -A nim00007

finetuning/generalists/util.py

@@ -78,7 +78,6 @@ def get_generalist_predictor(checkpoint, model_type, is_custom_model, return_sta
         )
 
 
-# TODO use model comparison func to generate the image data for qualitative comp
 def evaluate_checkpoint_for_dataset(
     checkpoint, model_type, dataset, experiment_folder,
     run_default_evaluation, run_amg, is_custom_model,