Update README.md

README.md

@@ -20,16 +20,16 @@ datasets:
 
 ⚖️ [**Model weights**](./MAESTRO_FLAIR-HUB_base/checkpoints/pretrain-epoch=99.ckpt) <br>
 ⚙️ [**Model configuration**](./MAESTRO_FLAIR-HUB_base/.hydra/config_resolved.yaml) <br>
-📂 [**Dataset splits**](https://huggingface.co/IGNF/MAESTRO_FLAIR-HUB_base/tree/main/dataset_splits) <br>
+📂 [**Dataset splits**](dataset_splits) <br>
 
 ## Abstract
 
-**MAESTRO** is a tailored adaptation of the Masked Autoencoder (MAE) framework that effectively orchestrates the use of multimodal, multitemporal, and multispectral Earth Observation (EO) data. Evaluated on four EO datasets, MAESTRO sets a new state-of-the-art on tasks that strongly rely on multitemporal dynamics, while remaining highly competitive on tasks dominated by a single monotemporal modality.
+**MAESTRO** is a tailored adaptation of the Masked Autoencoder (MAE) that effectively orchestrates the use of multimodal, multitemporal, and multispectral Earth Observation (EO) data. Evaluated on four EO datasets, MAESTRO sets a new state-of-the-art on tasks that strongly rely on multitemporal dynamics, while remaining competitive on tasks dominated by a single monotemporal modality.
 
 MAESTRO's contributions are as follows:
 - **Extensive benchmarking of multimodal and multitemporal SSL:** Impact evaluation of various fusion strategies for multimodal and multitemporal SSL.
 - **Patch-group-wise normalization:** Novel normalization scheme that normalizes reconstruction targets patch-wise within groups of highly correlated spectral bands.
-- **MAESTRO:** Novel adaptation of the MAE that combines optimized fusion strategies with our tailored patch-group-wise normalization..
+- **MAESTRO:** Novel adaptation of the MAE that combines optimized fusion strategies with patch-group-wise normalization.
 
 <div style="position: relative; text-align: center;">
     <img src="./media/Maestro_Overview.png" style="width: 100%; display: block; margin: 0 auto;"/>
@@ -59,7 +59,7 @@ We retain six distinct modalities:
 Below is the reconstruction loss during pre-training on the combined training and validation ensembles, using patch-group-wise normalization and modality-weighted averaging proportional to token counts.
 
 <div style="position: relative; text-align: center;">
-    <img src="./media/train_rec_loss_maeflair.png" style="width: 60%; display: block; margin: 0 auto;"/>
+    <img src="./media/Reconstruction_Loss.png" style="width: 100%; display: block; margin: 0 auto;"/>
 </div>
 
 <hr>
@@ -88,7 +88,7 @@ For optimal fine-tuning results with this model:
     - Patch size: 2
     - Channels: B02, B03, B04, B05, B06, B07, B08, B8A, B11, B12
 - Use fixed cross-dataset grids for positional encodings proportional to ground sampling distance: `grid_pos_enc` ≈ 1.6 * `crop_meters`
-  
+
 Note that modality names must match between pre-training and fine-tuning. 
 
 Below are cross-dataset evaluation results obtained with these guidelines on TreeSatAI-TS and PASTIS-HD.
@@ -110,7 +110,7 @@ Below are cross-dataset evaluation results obtained with these guidelines on Tre
 ## 🚀 Getting started
 
 Prerequisites: 
-- Fetch [Dataset splits](https://huggingface.co/IGNF/MAESTRO_FLAIR-HUB_base/tree/main/dataset_splits) and move them to each dataset directory
+- Fetch [Dataset splits](dataset_splits) and move them to each dataset directory
 - Fetch [model weights](./MAESTRO_FLAIR-HUB_base/checkpoints/pretrain-epoch=99.ckpt) and move them into `/path/to/experiments/MAESTRO_FLAIR-HUB_base/checkpoints/`
 - Fetch [model configuration](./MAESTRO_FLAIR-HUB_base/.hydra/config_resolved.yaml) and move it into `/path/to/experiments/MAESTRO_FLAIR-HUB_base/.hydra/`
 
@@ -131,7 +131,7 @@ poetry run python main.py \
   model.model=mae model.model_size=medium \
   model.fusion_mode=group model.inter_depth=3 \
   opt_pretrain.epochs=100 opt_probe.epochs=0 opt_finetune.epochs=0 \
-  opt_pretrain.batch_size=144 \
+  opt_pretrain.batch_size=9 trainer.num_nodes=4 \
   datasets.name_dataset=flair \
   datasets.flair.filter_inputs=[aerial,dem,spot,s2,s1_asc,s1_des] \
   datasets.flair.crop_meters=102.4 datasets.flair.grid_pos_enc=160 \
@@ -147,13 +147,13 @@ poetry run python main.py \
 
 Fine-tuning on TreeSatAI-TS:
 ```bash
-# batch size 32 on 1 node with 3 GPUs per node
+# batch size 24 on 1 node with 4 GPUs per node
 # load pre-trained model "MAESTRO_FLAIR-HUB_base"
 poetry run python main.py \
   model.model=mae model.model_size=medium \
   model.fusion_mode=group model.inter_depth=3 \
   opt_pretrain.epochs=0 opt_probe.epochs=10 opt_finetune.epochs=50 \
-  opt_probe.batch_size=96 opt_finetune.batch_size=96 \
+  opt_probe.batch_size=24 opt_finetune.batch_size=24 trainer.num_nodes=1 \
   opt_finetune.monitor=treesat_mlc_thresh/weighted_f1_val \
   datasets.name_dataset=treesatai_ts \
   datasets.treesatai_ts.filter_inputs=[aerial,s2,s1_asc,s1_des] \
@@ -169,13 +169,13 @@ poetry run python main.py \
 
 Fine-tuning on PASTIS-HD:
 ```bash
-# batch size 16 on 1 node with 3 GPUs per node
+# batch size 12 on 1 node with 4 GPUs per node
 # load pre-trained model "MAESTRO_FLAIR-HUB_base"
 poetry run python main.py \
   model.model=mae model.model_size=medium \
   model.fusion_mode=group model.inter_depth=3 \
   opt_pretrain.epochs=0 opt_probe.epochs=10 opt_finetune.epochs=50 \
-  opt_probe.batch_size=48 opt_finetune.batch_size=48 \
+  opt_probe.batch_size=12 opt_finetune.batch_size=12 trainer.num_nodes=1 \
   opt_finetune.monitor=pastis_seg/average_iou_val \
   datasets.name_dataset=pastis_hd \
   datasets.pastis_hd.filter_inputs=[spot,s2,s1_asc,s1_des] \
@@ -198,13 +198,14 @@ poetry run python main.py \
   model.model=mae model.model_size=medium \
   model.fusion_mode=group model.inter_depth=3 \
   opt_pretrain.epochs=0 opt_probe.epochs=15 opt_finetune.epochs=100 \
-  opt_probe.batch_size=48 opt_finetune.batch_size=48 \
+  opt_probe.batch_size=6 opt_finetune.batch_size=6 trainer.num_nodes=2 \
   opt_finetune.monitor=cosia/average_iou_val \
   datasets.name_dataset=flair \
   datasets.flair.version=flair2 \
-  datasets.flair.filter_inputs=[aerial,s2] \
+  datasets.flair.filter_inputs=[aerial,dem,s2] \
   datasets.flair.crop_meters=102.4 datasets.flair.grid_pos_enc=160 \
   datasets.flair.aerial.image_size=512 datasets.flair.aerial.patch_size.mae=16 \
+  datasets.flair.dem.image_size=512 datasets.flair.dem.patch_size.mae=32 \
   datasets.flair.s2.image_size=10 datasets.flair.s2.patch_size.mae=2 \
   datasets.root_dir=/path/to/dataset/dir datasets.flair.csv_dir=/path/to/dataset/dir/FLAIR-HUB datasets.flair.rel_dir=FLAIR-HUB \
   run.exp_dir=/path/to/experiments/dir run.exp_name=MAESTRO_FLAIR-HUB-x-FLAIR2_base \
@@ -219,12 +220,13 @@ poetry run python main.py \
   model.model=mae model.model_size=medium \
   model.fusion_mode=group model.inter_depth=3 \
   opt_pretrain.epochs=0 opt_probe.epochs=15 opt_finetune.epochs=100 \
-  opt_probe.batch_size=96 opt_finetune.batch_size=96 \
+  opt_probe.batch_size=6 opt_finetune.batch_size=6 trainer.num_nodes=4 \
   opt_finetune.monitor=cosia/average_iou_val \
   datasets.name_dataset=flair \
-  datasets.flair.filter_inputs=[aerial,s2,s1_asc,s1_des] \
+  datasets.flair.filter_inputs=[aerial,dem,s2,s1_asc,s1_des] \
   datasets.flair.crop_meters=102.4 datasets.flair.grid_pos_enc=160 \
   datasets.flair.aerial.image_size=512 datasets.flair.aerial.patch_size.mae=16 \
+  datasets.flair.dem.image_size=512 datasets.flair.dem.patch_size.mae=32 \
   datasets.flair.s2.image_size=10 datasets.flair.s2.patch_size.mae=2 \
   datasets.flair.s1_asc.image_size=10 datasets.flair.s1_asc.patch_size.mae=2 \
   datasets.flair.s1_des.image_size=10 datasets.flair.s1_des.patch_size.mae=2 \