Try pre-training Swin model + UNet decoder on Helios dataset

data/helios/v2_pastis/README.md

@@ -35,4 +35,8 @@ python -m rslp.main helios launch_finetune --helios_checkpoint_path /weka/dfive-
 
 # Helios with max pooling.
 python -m rslp.main helios launch_finetune --helios_checkpoint_path /weka/dfive-default/helios/checkpoints/joer/v0.1_base_latent_mim_space_time/step165000 --patch_size 8 --encoder_embedding_size 768 --image_name favyen/rslphelios2 --config_paths+=data/helios/v2_pastis/basecfg.yaml --config_paths+=data/helios/v2_pastis/basecfg_helios_ts.yaml --config_paths+=data/helios/v2_shared/helios_freeze_then_lowlr.yaml --config_paths+=data/helios/v2_shared/helios_ts_simple_maxpool_sentinel2.yaml --cluster+=ai2/ceres-cirrascale --cluster+=ai2/saturn-cirrascale --rslp_project 2025_06_06_helios_finetuning --experiment_id v2_pastis_ts_maxpool_helios_latent_mim_space_time
+
+# sp
+python -m rslp.main common beaker_train --config_paths+=data/helios/v2_pastis/basecfg.yaml --config_paths+=data/helios/v2_pastis/basecfg_sp.yaml --config_paths+=data/helios/v2_shared/sp.yaml --project_id 2025_06_06_helios_finetuning --experiment_id v2_pastis_uni_sp --cluster+=ai2/ceres --cluster+=ai2/jupiter '--weka_mounts+={"bucket_name":"dfive-default","mount_path":"/weka/dfive-default"}' --image_name favyen/rslphelios13
+python -m rslp.main common beaker_train --config_paths+=data/helios/v2_pastis/basecfg.yaml --config_paths+=data/helios/v2_pastis/basecfg_sp_ts.yaml --config_paths+=data/helios/v2_shared/sp.yaml --project_id 2025_06_06_helios_finetuning --experiment_id v2_pastis_ts_sp --cluster+=ai2/ceres --cluster+=ai2/jupiter '--weka_mounts+={"bucket_name":"dfive-default","mount_path":"/weka/dfive-default"}' --image_name favyen/rslphelios13
 ```

data/helios/v2_pastis/basecfg_sp.yaml

@@ -0,0 +1,62 @@
+model:
+  init_args:
+    model:
+      init_args:
+        encoder:
+          - class_path: rslp.swin_pretrain.model.Model
+        decoders:
+          segment:
+            - class_path: rslearn.models.conv.Conv
+              init_args:
+                in_channels: 128
+                out_channels: 20
+                kernel_size: 1
+                activation:
+                  class_path: torch.nn.Identity
+            - class_path: rslearn.models.pick_features.PickFeatures
+              init_args:
+                indexes: [0]
+                collapse: true
+            - class_path: rslearn.train.tasks.segmentation.SegmentationHead
+data:
+  init_args:
+    inputs:
+      image:
+        data_type: "raster"
+        layers: ["sentinel2"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      targets:
+        data_type: "raster"
+        layers: ["label"]
+        bands: ["class"]
+        is_target: true
+    default_config:
+      transforms:
+        - class_path: rslearn.train.transforms.concatenate.Concatenate
+          init_args:
+            selections:
+              # PASTIS is missing B01 and B09.
+              # We use B02 to fill in B01 and B8A to fill in B09.
+              image: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+            output_selector: image
+        - class_path: rslearn.train.transforms.normalize.Normalize
+          init_args:
+            mean: 0
+            std: 10000
+    train_config:
+      transforms:
+        - class_path: rslearn.train.transforms.concatenate.Concatenate
+          init_args:
+            selections:
+              # PASTIS is missing B01 and B09.
+              # We use B02 to fill in B01 and B8A to fill in B09.
+              image: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+            output_selector: image
+        - class_path: rslearn.train.transforms.normalize.Normalize
+          init_args:
+            mean: 0
+            std: 10000
+        - class_path: rslearn.train.transforms.flip.Flip
+          init_args:
+            image_selectors: ["image", "target/segment/classes", "target/segment/valid"]

data/helios/v2_pastis/basecfg_sp_ts.yaml

@@ -0,0 +1,139 @@
+model:
+  init_args:
+    model:
+      init_args:
+        encoder:
+          - class_path: rslp.swin_pretrain.model.Model
+        decoders:
+          segment:
+            - class_path: rslearn.models.conv.Conv
+              init_args:
+                in_channels: 128
+                out_channels: 20
+                kernel_size: 1
+                activation:
+                  class_path: torch.nn.Identity
+            - class_path: rslearn.models.pick_features.PickFeatures
+              init_args:
+                indexes: [0]
+                collapse: true
+            - class_path: rslearn.train.tasks.segmentation.SegmentationHead
+data:
+  init_args:
+    inputs:
+      sentinel2_0:
+        data_type: "raster"
+        layers: ["sentinel2"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      sentinel2_1:
+        data_type: "raster"
+        layers: ["sentinel2.1"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      sentinel2_2:
+        data_type: "raster"
+        layers: ["sentinel2.2"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      sentinel2_3:
+        data_type: "raster"
+        layers: ["sentinel2.3"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      sentinel2_4:
+        data_type: "raster"
+        layers: ["sentinel2.4"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      sentinel2_5:
+        data_type: "raster"
+        layers: ["sentinel2.5"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      sentinel2_6:
+        data_type: "raster"
+        layers: ["sentinel2.6"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      sentinel2_7:
+        data_type: "raster"
+        layers: ["sentinel2.7"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      sentinel2_8:
+        data_type: "raster"
+        layers: ["sentinel2.8"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      sentinel2_9:
+        data_type: "raster"
+        layers: ["sentinel2.9"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      sentinel2_10:
+        data_type: "raster"
+        layers: ["sentinel2.10"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      sentinel2_11:
+        data_type: "raster"
+        layers: ["sentinel2.11"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12"]
+        passthrough: true
+      targets:
+        data_type: "raster"
+        layers: ["label"]
+        bands: ["class"]
+        is_target: true
+    default_config:
+      transforms:
+        - class_path: rslearn.train.transforms.concatenate.Concatenate
+          init_args:
+            selections:
+              # PASTIS is missing B01 and B09.
+              # We use B02 to fill in B01 and B8A to fill in B09.
+              sentinel2_0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_1: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_2: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_3: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_4: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_5: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_6: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_7: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_8: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_9: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_10: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_11: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+            output_selector: image
+        - class_path: rslearn.train.transforms.normalize.Normalize
+          init_args:
+            mean: 0
+            std: 10000
+    train_config:
+      transforms:
+        - class_path: rslearn.train.transforms.concatenate.Concatenate
+          init_args:
+            selections:
+              # PASTIS is missing B01 and B09.
+              # We use B02 to fill in B01 and B8A to fill in B09.
+              sentinel2_0: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_1: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_2: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_3: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_4: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_5: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_6: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_7: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_8: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_9: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_10: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+              sentinel2_11: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 7]
+            output_selector: image
+        - class_path: rslearn.train.transforms.normalize.Normalize
+          init_args:
+            mean: 0
+            std: 10000
+        - class_path: rslearn.train.transforms.flip.Flip
+          init_args:
+            image_selectors: ["image", "target/segment/classes", "target/segment/valid"]

data/helios/v2_satlas_marine_infra_128/README.md

@@ -35,4 +35,8 @@ python -m rslp.main helios launch_finetune --helios_checkpoint_path /weka/dfive-
 
 # Helios with max pooling.
 python -m rslp.main helios launch_finetune --helios_checkpoint_path /weka/dfive-default/helios/checkpoints/joer/v0.1_base_latent_mim_space_time/step165000 --patch_size 8 --encoder_embedding_size 768 --image_name favyen/rslphelios2 --config_paths+=data/helios/v2_satlas_marine_infra_128/basecfg.yaml --config_paths+=data/helios/v2_satlas_marine_infra_128/basecfg_helios_ts.yaml --config_paths+=data/helios/v2_shared/helios_freeze_then_lowlr.yaml --config_paths+=data/helios/v2_shared/helios_ts_simple_maxpool_sentinel2.yaml --cluster+=ai2/ceres-cirrascale --cluster+=ai2/saturn-cirrascale --rslp_project 2025_06_06_helios_finetuning --experiment_id v2_satlas_marine_infra_128_ts_maxpool_helios_latent_mim_space_time
+
+# sp
+python -m rslp.main common beaker_train --config_paths+=data/helios/v2_satlas_marine_infra_128/basecfg.yaml --config_paths+=data/helios/v2_satlas_marine_infra_128/basecfg_sp.yaml --config_paths+=data/helios/v2_shared/sp.yaml --project_id 2025_06_06_helios_finetuning --experiment_id v2_satlas_marine_infra_128_uni_sp --cluster+=ai2/ceres --cluster+=ai2/jupiter '--weka_mounts+={"bucket_name":"dfive-default","mount_path":"/weka/dfive-default"}' --image_name favyen/rslphelios13
+python -m rslp.main common beaker_train --config_paths+=data/helios/v2_satlas_marine_infra_128/basecfg.yaml --config_paths+=data/helios/v2_satlas_marine_infra_128/basecfg_sp_ts.yaml --config_paths+=data/helios/v2_shared/sp.yaml --project_id 2025_06_06_helios_finetuning --experiment_id v2_satlas_marine_infra_128_ts_sp --cluster+=ai2/ceres --cluster+=ai2/jupiter '--weka_mounts+={"bucket_name":"dfive-default","mount_path":"/weka/dfive-default"}' --image_name favyen/rslphelios13
 ```

data/helios/v2_satlas_marine_infra_128/basecfg_sp.yaml

@@ -0,0 +1,58 @@
+model:
+  init_args:
+    model:
+      init_args:
+        encoder:
+          - class_path: rslp.swin_pretrain.model.Model
+            init_args:
+              target_resolution_factor: null
+          - class_path: rslearn.models.fpn.Fpn
+            init_args:
+              in_channels: [128, 256, 512, 1024]
+              out_channels: 128
+        decoders:
+          detect:
+            - class_path: rslearn.models.faster_rcnn.FasterRCNN
+              init_args:
+                downsample_factors: [4, 8, 16, 32]
+                num_channels: 128
+                num_classes: 3
+                anchor_sizes: [[32], [64], [128], [256]]
+data:
+  init_args:
+    inputs:
+      image:
+        data_type: "raster"
+        layers: ["sentinel2"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12", "B01", "B09"]
+        passthrough: true
+        dtype: FLOAT32
+      mask:
+        data_type: "raster"
+        layers: ["mask"]
+        bands: ["mask"]
+        passthrough: true
+        dtype: FLOAT32
+        is_target: true
+      targets:
+        data_type: "vector"
+        layers: ["label"]
+        is_target: true
+    default_config:
+      transforms:
+        - class_path: rslearn.train.transforms.normalize.Normalize
+          init_args:
+            mean: 0
+            std: 10000
+        - class_path: rslp.transforms.mask.Mask
+    train_config:
+      transforms:
+        - class_path: rslearn.train.transforms.normalize.Normalize
+          init_args:
+            mean: 0
+            std: 10000
+        - class_path: rslp.transforms.mask.Mask
+        - class_path: rslearn.train.transforms.flip.Flip
+          init_args:
+            image_selectors: ["image"]
+            box_selectors: ["target/detect"]

data/helios/v2_satlas_marine_infra_128/basecfg_sp_ts.yaml

@@ -0,0 +1,92 @@
+model:
+  init_args:
+    model:
+      init_args:
+        encoder:
+          - class_path: rslp.swin_pretrain.model.Model
+            init_args:
+              target_resolution_factor: null
+          - class_path: rslearn.models.fpn.Fpn
+            init_args:
+              in_channels: [128, 256, 512, 1024]
+              out_channels: 128
+        decoders:
+          detect:
+            - class_path: rslearn.models.faster_rcnn.FasterRCNN
+              init_args:
+                downsample_factors: [4, 8, 16, 32]
+                num_channels: 128
+                num_classes: 3
+                anchor_sizes: [[32], [64], [128], [256]]
+data:
+  init_args:
+    inputs:
+      image1:
+        data_type: "raster"
+        layers: ["sentinel2"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12", "B01", "B09"]
+        passthrough: true
+        dtype: FLOAT32
+      image2:
+        data_type: "raster"
+        layers: ["sentinel2.1"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12", "B01", "B09"]
+        passthrough: true
+        dtype: FLOAT32
+      image3:
+        data_type: "raster"
+        layers: ["sentinel2.2"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12", "B01", "B09"]
+        passthrough: true
+        dtype: FLOAT32
+      image4:
+        data_type: "raster"
+        layers: ["sentinel2.3"]
+        bands: ["B02", "B03", "B04", "B08", "B05", "B06", "B07", "B8A", "B11", "B12", "B01", "B09"]
+        passthrough: true
+        dtype: FLOAT32
+      mask:
+        data_type: "raster"
+        layers: ["mask"]
+        bands: ["mask"]
+        passthrough: true
+        dtype: FLOAT32
+        is_target: true
+      targets:
+        data_type: "vector"
+        layers: ["label"]
+        is_target: true
+    default_config:
+      transforms:
+        - class_path: rslearn.train.transforms.concatenate.Concatenate
+          init_args:
+            selections:
+              image1: []
+              image2: []
+              image3: []
+              image4: []
+            output_selector: image
+        - class_path: rslearn.train.transforms.normalize.Normalize
+          init_args:
+            mean: 0
+            std: 10000
+        - class_path: rslp.transforms.mask.Mask
+    train_config:
+      transforms:
+        - class_path: rslearn.train.transforms.concatenate.Concatenate
+          init_args:
+            selections:
+              image1: []
+              image2: []
+              image3: []
+              image4: []
+            output_selector: image
+        - class_path: rslearn.train.transforms.normalize.Normalize
+          init_args:
+            mean: 0
+            std: 10000
+        - class_path: rslp.transforms.mask.Mask
+        - class_path: rslearn.train.transforms.flip.Flip
+          init_args:
+            image_selectors: ["image"]
+            box_selectors: ["target/detect"]

data/helios/v2_satlas_solar_farm_128/README.md

@@ -35,4 +35,8 @@ python -m rslp.main helios launch_finetune --helios_checkpoint_path /weka/dfive-
 
 # Helios with max pooling.
 python -m rslp.main helios launch_finetune --helios_checkpoint_path /weka/dfive-default/helios/checkpoints/joer/v0.1_base_latent_mim_space_time/step165000 --patch_size 8 --encoder_embedding_size 768 --image_name favyen/rslphelios2 --config_paths+=data/helios/v2_satlas_solar_farm_128/basecfg.yaml --config_paths+=data/helios/v2_satlas_solar_farm_128/basecfg_helios_ts.yaml --config_paths+=data/helios/v2_shared/helios_freeze_then_lowlr.yaml --config_paths+=data/helios/v2_shared/helios_ts_simple_maxpool_sentinel2.yaml --cluster+=ai2/ceres-cirrascale --cluster+=ai2/saturn-cirrascale --rslp_project 2025_06_06_helios_finetuning --experiment_id v2_satlas_solar_farm_128_ts_maxpool_helios_latent_mim_space_time
+
+# sp
+python -m rslp.main common beaker_train --config_paths+=data/helios/v2_satlas_solar_farm_128/basecfg.yaml --config_paths+=data/helios/v2_satlas_solar_farm_128/basecfg_sp.yaml --config_paths+=data/helios/v2_shared/sp.yaml --project_id 2025_06_06_helios_finetuning --experiment_id v2_satlas_solar_farm_128_uni_sp --cluster+=ai2/ceres --cluster+=ai2/jupiter '--weka_mounts+={"bucket_name":"dfive-default","mount_path":"/weka/dfive-default"}' --image_name favyen/rslphelios13
+python -m rslp.main common beaker_train --config_paths+=data/helios/v2_satlas_solar_farm_128/basecfg.yaml --config_paths+=data/helios/v2_satlas_solar_farm_128/basecfg_sp_ts.yaml --config_paths+=data/helios/v2_shared/sp.yaml --project_id 2025_06_06_helios_finetuning --experiment_id v2_satlas_solar_farm_128_ts_sp --cluster+=ai2/ceres --cluster+=ai2/jupiter '--weka_mounts+={"bucket_name":"dfive-default","mount_path":"/weka/dfive-default"}' --image_name favyen/rslphelios13
 ```