Merge pull request #544 from allenai/hadriens/burnrisk_task

New eval task: 8 day burn-risk forecast

Author: Farbum

olmoearth_pretrain/evals/studio_ingest/registry.json

@@ -115,6 +115,115 @@
       "weka_path": "/weka/dfive-default/olmoearth/eval_datasets/awf_lulc_map",
       "window_size": 64
     },
+    "burnrisk_8d_nbac": {
+      "classes": [
+        "0",
+        "1"
+      ],
+      "imputes": [],
+      "is_multilabel": false,
+      "modalities": [
+        "sentinel1",
+        "sentinel2_l2a"
+      ],
+      "name": "burnrisk_8d_nbac",
+      "norm_stats": {
+        "sentinel2_l2a": {
+          "B01": {
+            "max": 19963.0,
+            "mean": 2797.991553441403,
+            "min": 0.0,
+            "std": 3462.128374048374
+          },
+          "B02": {
+            "max": 17056.0,
+            "mean": 2742.4601498804145,
+            "min": 0.0,
+            "std": 3357.983085064015
+          },
+          "B03": {
+            "max": 17152.0,
+            "mean": 2766.9766015413234,
+            "min": 0.0,
+            "std": 3191.070943297956
+          },
+          "B04": {
+            "max": 17184.0,
+            "mean": 2720.677049800691,
+            "min": 0.0,
+            "std": 3164.727976467343
+          },
+          "B05": {
+            "max": 16850.0,
+            "mean": 3030.2717098060057,
+            "min": 0.0,
+            "std": 3132.0173541691083
+          },
+          "B06": {
+            "max": 16603.0,
+            "mean": 3382.8399415360086,
+            "min": 0.0,
+            "std": 2829.0052701581553
+          },
+          "B07": {
+            "max": 16437.0,
+            "mean": 3437.354116290194,
+            "min": 0.0,
+            "std": 2689.959382813977
+          },
+          "B08": {
+            "max": 16296.0,
+            "mean": 3559.7809055540793,
+            "min": 0.0,
+            "std": 2747.156814362987
+          },
+          "B09": {
+            "max": 17383.0,
+            "mean": 4172.207622110019,
+            "min": 0.0,
+            "std": 3629.4698043631947
+          },
+          "B11": {
+            "max": 15184.0,
+            "mean": 1905.4351294180176,
+            "min": 0.0,
+            "std": 1490.2586203959572
+          },
+          "B12": {
+            "max": 15118.0,
+            "mean": 1473.891998086633,
+            "min": 0.0,
+            "std": 1316.966842986169
+          },
+          "B8A": {
+            "max": 16164.0,
+            "mean": 3479.5087070954028,
+            "min": 0.0,
+            "std": 2560.063675143769
+          }
+        }
+      },
+      "num_classes": 2,
+      "num_timesteps": 8,
+      "source_path": "/weka/dfive-default/rslearn-eai/datasets/wildfire/canada_nbac",
+      "split_stats": {
+        "test": {
+          "count": 1000
+        },
+        "train": {
+          "count": 1000
+        },
+        "val": {
+          "count": 500
+        }
+      },
+      "split_tag_key": "eval_split",
+      "task_type": "segmentation",
+      "timeseries": true,
+      "use_pretrain_norm": true,
+      "weka_path": "/weka/dfive-default/olmoearth/eval_datasets/burnrisk_8d_nbac",
+      "window_size": 50
+    },
     "canada_wildfire_sat_eval_split": {
       "classes": [
         "0",
@@ -1142,6 +1251,6 @@
       "window_size": 16
     }
   },
-  "updated_at": "2026-02-24T00:15:34.513906",
+  "updated_at": "2026-04-15T23:29:08.666694",
   "version": "1.0"
 }

olmoearth_pretrain/internal/all_evals.py

@@ -278,16 +278,16 @@ def load_user_module(path: str) -> Any:
         epochs=50,
         eval_mode=EvalMode.LINEAR_PROBE,
     ),
-    "canada_wildfire_sat_eval_split": DownstreamTaskConfig(
-        dataset="canada_wildfire_sat_eval_split",
+    "burnrisk_8d_nbac": DownstreamTaskConfig(
+        dataset="burnrisk_8d_nbac",
         embedding_batch_size=32,
         probe_batch_size=16,
-        patch_size=5,  # TODO: This is changeable but we should know the valid sizes for inputs
-        num_workers=2,
+        patch_size=5,
+        num_workers=4,
         pooling_type=PoolingType.MEAN,
         norm_stats_from_pretrained=True,
         norm_method=NormMethod.NORM_NO_CLIP_2_STD,
-        probe_lr=0.1,
+        probe_lr=0.0001,
         eval_interval=Duration.epochs(10),
         input_modalities=[Modality.SENTINEL2_L2A.name],
         epochs=50,

scripts/base_loop_evals/script.py

@@ -1,6 +1,6 @@
 """Same as official script except EVAL_TASKS.
 
-Only tolbi_crop, canada_wildfire_sat_eval_split, yemen_crop,
+Only tolbi_crop, burnrisk_8d_nbac, yemen_crop,
 geo_ecosystem_annual_test, forest_loss_driver, nigeria_settlement,
 nandi_crop_map, awf_lulc_map, oil_spill_detection as loop evals every 5k steps.
 """
@@ -29,6 +29,7 @@
 from olmoearth_pretrain.data.dataloader import OlmoEarthDataLoaderConfig
 from olmoearth_pretrain.data.dataset import OlmoEarthDatasetConfig
 from olmoearth_pretrain.evals.datasets.normalize import NormMethod
+from olmoearth_pretrain.evals.metrics import EvalMetric
 from olmoearth_pretrain.internal.common import (
     build_common_components as build_common_components_default,
 )
@@ -103,21 +104,23 @@
         epochs=50,
         eval_mode=EvalMode.LINEAR_PROBE,
     ),
-    "canada_wildfire_sat_eval_split": DownstreamTaskConfig(
-        dataset="canada_wildfire_sat_eval_split",
+    "burnrisk_8d_nbac": DownstreamTaskConfig(
+        dataset="burnrisk_8d_nbac",
         embedding_batch_size=32,
         probe_batch_size=16,
         patch_size=5,
-        num_workers=2,
+        num_workers=4,
         pooling_type=PoolingType.MEAN,
         norm_stats_from_pretrained=True,
         norm_method=NormMethod.NORM_NO_CLIP_2_STD,
-        probe_lr=0.1,
+        probe_lr=0.0001,
         eval_interval=LOOP_EVAL_INTERVAL,
         input_modalities=[Modality.SENTINEL2_L2A.name],
         epochs=50,
         eval_mode=EvalMode.LINEAR_PROBE,
         use_dice_loss=True,
+        primary_metric=EvalMetric.CLASS_F1,
+        primary_metric_class=1,
     ),
     "yemen_crop": DownstreamTaskConfig(
         dataset="yemen_crop",