Configuration template for Swin

Signed-off-by: João Lucas de Sousa Almeida <[email protected]>

Author: Joao-L-S-Almeida

examples/confs/sen1floods11_swin.yaml

@@ -0,0 +1,144 @@
+# lightning.pytorch==2.1.1
+seed_everything: 0
+# save the model checkpoints to a separate location then the logger
+# ModelCheckpoint:
+#   dirpath: /dccstor/geofm-finetuning/carlosgomes/torchgeo_floods/sen1floods11_terratorch
+trainer:
+  accelerator: auto
+  strategy: auto
+  devices: auto
+  num_nodes: 1
+  precision: 16-mixed
+  logger:
+    class_path: lightning.pytorch.loggers.mlflow.MLFlowLogger
+    init_args:
+      experiment_name: "test_experiment_mlflow"
+      run_name: "plot_every_2_epochs"
+      save_dir: <your_path_here>/mlflow
+  callbacks:
+    - class_path: RichProgressBar
+    - class_path: LearningRateMonitor
+      init_args:
+        logging_interval: epoch
+    - class_path: EarlyStopping
+      init_args:
+        monitor: val/loss
+        patience: 20
+
+  max_epochs: 200
+  check_val_every_n_epoch: 1
+  log_every_n_steps: 50
+  enable_checkpointing: true
+  default_root_dir: <your_path_here>/torchgeo_floods
+data:
+  class_path: GenericNonGeoSegmentationDataModule
+  init_args:
+    batch_size: 4
+    num_workers: 8
+    constant_scale: 0.0001
+    dataset_bands:
+      - RED
+      - GREEN
+      - BLUE
+      - NIR_NARROW
+      - SWIR_1
+      - SWIR_2
+    output_bands:
+      - BLUE
+      - GREEN
+      - RED
+      - NIR_NARROW
+      - SWIR_1
+      - SWIR_2
+    rgb_indices:
+      - 2
+      - 1
+      - 0
+    train_data_root: <sen1floods11_root>/v1.1/data/flood_events/HandLabeled/S2Hand/
+    train_label_data_root: <sen1floods11_root>/v1.1/data/flood_events/HandLabeled/LabelHand
+    val_data_root: <sen1floods11_root>/v1.1/data/flood_events/HandLabeled/S2Hand/
+    val_label_data_root: <sen1floods11_root>/v1.1/data/flood_events/HandLabeled/LabelHand
+    test_data_root: <sen1floods11_root>/v1.1/data/flood_events/HandLabeled/S2Hand/
+    test_label_data_root: <sen1floods11_root>/v1.1/data/flood_events/HandLabeled/LabelHand
+    # these must be obtained by running terratorch/examples/scripts/convert_sen1floods11_splits.py on the original split csv files
+    train_split: <sen1floods11_root>/splits/splits/flood_handlabeled/flood_train_data.txt
+    test_split: <sen1floods11_root>/splits/splits/flood_handlabeled/flood_test_data.txt
+    val_split: <sen1floods11_root>/splits/splits/flood_handlabeled/flood_valid_data.txt
+    img_grep: "*_S2GeodnHand.tif"
+    label_grep: "*_LabelHand.tif"
+    no_label_replace: -1
+    no_data_replace: 0
+    means:
+      - 0.107582
+      - 0.13471393
+      - 0.12520133
+      - 0.3236181
+      - 0.2341743
+      - 0.15878009
+    stds:
+      - 0.07145836
+      - 0.06783548
+      - 0.07323416
+      - 0.09489725
+      - 0.07938496
+      - 0.07089546
+    num_classes: 2
+
+model:
+  class_path: terratorch.tasks.SemanticSegmentationTask
+  init_args:
+    model_args:
+      decoder: UperNetDecoder
+      pretrained: true
+      backbone: prithvi_swin_B
+      backbone_drop_path_rate: 0.3
+      backbone_window_size: 7
+      decoder_channels: 256
+      in_channels: 6
+      bands:
+        - BLUE
+        - GREEN
+        - RED
+        - NIR_NARROW
+        - SWIR_1
+        - SWIR_2
+      num_frames: 1
+      num_classes: 2
+      head_dropout: 0.1
+      head_channel_list:
+        - 256
+    loss: ce
+    plot_on_val: 2
+    # aux_heads:
+    #   - name: aux_head
+    #     decoder: FCNDecoder
+    #     decoder_args:
+    #       decoder_channels: 256
+    #       decoder_in_index: 2
+    #       decoder_num_convs: 1
+    #       head_channel_list:
+    #         - 64
+    # aux_loss:
+    #   aux_head: 1.0
+    ignore_index: -1
+    class_weights:
+      - 0.3
+      - 0.7
+    freeze_backbone: false
+    freeze_decoder: false
+    model_factory: PrithviModelFactory
+    # tiled_inference_parameters:
+    #   h_crop: 512
+    #   h_stride: 512
+    #   w_crop: 512
+    #   w_stride: 512
+    #   average_patches: true
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 6.e-5
+    weight_decay: 0.05
+lr_scheduler:
+  class_path: ReduceLROnPlateau
+  init_args:
+    monitor: val/loss