add finetune flag

Author: yawenzzzz

olmoearth_pretrain/internal/all_evals.py

@@ -80,6 +80,7 @@ def load_user_module(path: str) -> Any:
         num_workers=4,
         pooling_type=PoolingType.MEAN,
         norm_stats_from_pretrained=False,
+        norm_method=NormMethod.NORM_NO_CLIP_2_STD,
         eval_interval=Duration.epochs(5),
         eval_mode=EvalMode.KNN,
     ),
@@ -117,6 +118,7 @@ def load_user_module(path: str) -> Any:
         num_workers=2,
         pooling_type=PoolingType.MEAN,
         norm_stats_from_pretrained=False,
+        norm_method=NormMethod.NORM_NO_CLIP_2_STD,
         probe_lr=0.1,
         eval_interval=Duration.epochs(10),
         eval_mode=EvalMode.LINEAR_PROBE,
@@ -127,7 +129,8 @@ def load_user_module(path: str) -> Any:
         probe_batch_size=8,
         num_workers=2,
         pooling_type=PoolingType.MEAN,
-        norm_stats_from_pretrained=True,
+        norm_stats_from_pretrained=False,
+        norm_method=NormMethod.NORM_NO_CLIP_2_STD,
         probe_lr=0.1,
         eval_interval=Duration.epochs(10),
         eval_mode=EvalMode.LINEAR_PROBE,
@@ -139,6 +142,7 @@ def load_user_module(path: str) -> Any:
         num_workers=8,
         pooling_type=PoolingType.MEAN,
         norm_stats_from_pretrained=False,
+        norm_method=NormMethod.NORM_NO_CLIP_2_STD,
         probe_lr=0.01,
         eval_interval=Duration.epochs(10),
         eval_mode=EvalMode.LINEAR_PROBE,
@@ -149,7 +153,7 @@ def load_user_module(path: str) -> Any:
         probe_batch_size=128,
         num_workers=4,
         pooling_type=PoolingType.MEAN,
-        norm_stats_from_pretrained=False,
+        norm_stats_from_pretrained=True,
         probe_lr=0.1,
         eval_interval=Duration.epochs(10),
         eval_mode=EvalMode.LINEAR_PROBE,

olmoearth_pretrain/internal/full_eval_sweep_finetune.py

@@ -1,4 +1,4 @@
-"""Launch fine-tune evaluation sweeps for OlmoEarth Pretrain checkpoints.
+"""Launch fine-tune evaluation sweeps for OlmoEarth and other models.
 
 Example run:
 python olmoearth_pretrain/internal/full_eval_sweep_finetune.py --project_name 2025_10_08_phase2_finetune --module_path olmoearth_pretrain/evals/models/clay/clay_launch.py --cluster ai2/titan --model_name clay --clay --defaults_only

scripts/get_max_eval_metrics_from_wandb.py

@@ -12,11 +12,11 @@
     MODELS_WITH_MULTIPLE_SIZES,
     BaselineModelName,
 )
-from olmoearth_pretrain.internal.all_evals import EVAL_TASKS
+from olmoearth_pretrain.internal.all_evals import EVAL_TASKS, FT_EVAL_TASKS
 from olmoearth_pretrain.train.callbacks.evaluator_callback import EvalMode
 
 WANDB_ENTITY = "eai-ai2"
-METRICS = EVAL_TASKS.keys()
+METRICS = list(EVAL_TASKS.keys())
 
 # Dataset partitions to consider (excluding default)
 PARTITIONS = [
@@ -379,9 +379,18 @@ def save_metrics_to_csv(metrics_dict: dict[str, dict[str, float]], filename: str
         action="store_true",
         help="Report test metrics based on the configuration of the validation results witht the highest score",
     )
+    parser.add_argument(
+        "--finetune",
+        action="store_true",
+        help="Use finetune evaluation tasks when determining metrics",
+    )
 
     args = parser.parse_args()
 
+    global METRICS
+    selected_tasks = FT_EVAL_TASKS if args.finetune else EVAL_TASKS
+    METRICS = list(selected_tasks.keys())
+
     if args.per_partition:
         if not args.run_prefix:
             parser.error("--per-partition requires run_prefix to be specified")