Soft contrastive loss, scalable alignment training implementation

configs/callbacks/batch_size_finder.yaml

@@ -1,5 +1,6 @@
 # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.BatchSizeFinder.html
 
 _target_: lightning.pytorch.callbacks.BatchSizeFinder
-mode: 'binsearch'
-init_val: 16
+mode: "power"
+init_val: 16
+max_val: 1024

configs/callbacks/early_stopping.yaml

@@ -5,7 +5,7 @@ early_stopping:
   monitor: ??? # quantity to be monitored, must be specified !!!
   min_delta: 0. # minimum change in the monitored quantity to qualify as an improvement
   patience: 3 # number of checks with no improvement after which training will be stopped
-  verbose: False # verbosity mode
+  verbose: true # verbosity mode
   mode: "min" # "max" means higher metric value is better, can be also "min"
   strict: True # whether to crash the training if monitor is not found in the validation metrics
   check_finite: True # when set True, stops training when the monitor becomes NaN or infinite

configs/data/butterfly_coords.yaml

@@ -16,5 +16,5 @@ num_workers: 8
 pin_memory: true
 split_mode: "from_file"
 save_split: false
-saved_split_file_name: "split_indices_s2bms_2024-08-14-1459.pth"
+saved_split_file_name: "s2bms_unlabelled_union_val_test.pth"
 seed: ${seed}

configs/data/butterfly_coords_text.yaml

@@ -24,5 +24,5 @@ num_workers: 8
 pin_memory: true
 split_mode: "from_file"
 save_split: false
-saved_split_file_name: "split_indices_s2bms_2024-08-14-1459.pth"
+saved_split_file_name: "s2bms_unlabelled_union_val_test.pth"
 seed: ${seed}

configs/data/butterfly_coords_text_unlabelled.yaml

@@ -24,5 +24,5 @@ num_workers: 8
 pin_memory: true
 split_mode: "from_file"
 save_split: false
-saved_split_file_name: "split_indices_s2bms+s2bms-unlabelled-20260529_2026-05-29-1438.pth"
+saved_split_file_name: "s2bms_unlabelled_union_val_test.pth"
 seed: ${seed}

configs/data/butterfly_full_param_example.yaml

@@ -32,5 +32,5 @@ num_workers: 8
 pin_memory: true
 split_mode: "from_file"
 save_split: false
-saved_split_file_name: "split_indices_s2bms_2024-08-14-1459.pth"
+saved_split_file_name: "s2bms_unlabelled_union_val_test.pth"
 seed: ${seed}

configs/data/butterfly_s2_rgb.yaml

@@ -21,5 +21,5 @@ num_workers: 1
 pin_memory: false
 split_mode: "from_file" # default 'random'
 save_split: false
-saved_split_file_name: "split_indices_s2bms_2024-08-14-1459.pth"
+saved_split_file_name: "s2bms_unlabelled_union_val_test.pth"
 seed: ${seed}

configs/experiment/alignment.yaml

@@ -4,13 +4,48 @@
 # python train.py experiment=example
 
 defaults:
-  - override /model: geoclip_alignment
-  - override /data: butterfly_coords_text
-  - override /metrics: contrastive_similarities
+  - override /model: s2bms_alignment
+  - override /data: s2bms_prediction
+  - override /metrics: s2bms_alignment
 
 # all parameters below will be merged with parameters from default configurations set above
 # this allows you to overwrite only specified parameters
 
+model:
+  geo_encoder:
+    _target_: src.models.components.geo_encoders.geoclip.GeoClipCoordinateEncoder
+    geo_data_name: coords
+  geo_adapter:
+    _target_: src.models.components.projectors_adapters.mlp_projector.MLPProjector
+    nn_layers: 2
+    hidden_dim: 512
+    output_dim: 768
+  text_encoder:
+    _target_: src.models.components.text_encoders.clip_text_encoder.ClipTextEncoder
+    hf_cache_dir: ${paths.huggingface_cache}
+    use_geoclip_projector: false
+  loss_fn:
+    _target_: src.models.components.loss_fns.clip_loss.ClipLoss
+    temperature: 0.07
+  trainable_modules: [loss_fn.log_temp]
+
+data:
+  dataset:
+    modalities:
+      coords:
+    use_target_data: false
+    use_aux_data: "all"
+    dtype: float32
+  caption_builder:
+    _target_: src.data.butterfly_caption_builder.ButterflyCaptionBuilder
+    templates_fname: v5.json
+    concepts_fname: v4.json
+    data_dir: ${paths.data_dir}/s2bms
+    seed: ${seed}
+    n_captions_for_validation: 50
+  saved_split_file_name: "s2bms_unlabelled_union_val_test.pth"
+  batch_size: 32
+
 tags: ["alignment", "geoclip_coords", "geoclip_text"]
 
 seed: 12345
@@ -19,9 +54,6 @@ trainer:
   min_epochs: 10
   max_epochs: 100
 
-data:
-  batch_size: 64
-
 logger:
   wandb:
     tags: ${tags}

configs/experiment/alignment_template.yaml

@@ -0,0 +1,95 @@
+# @package _global_
+
+# to execute this experiment run:
+# python train.py experiment=example
+
+# TODO: description
+
+defaults:
+  - override /model: s2bms_alignment
+  - override /data: s2bms_prediction
+  - override /metrics: s2bms_alignment
+
+tags:
+  - alignment
+  - aef
+  - aef_${str:${data.dataset.modalities.tessera_avr.size}} #TODO modality
+  #  - aef_${str:${data.dataset.modalities.tessera.size}} #TODO modality
+  - ${str:${seed}}
+
+# (unlab_)(geo_encoder)_(eo_mod)_(geo_adapter)_(text_encoder)_(text_adapter)_(batch_size) ,_ generated by_default plus extra
+#experiment_name: unlab_avr_aef_128_mlp_cliptext # TODO: to fully over-write the name
+#experiment_name_extra: # TODO any extra unique features about the experiment
+
+projects: s2bms_alignment
+
+model:
+  geo_encoder:
+    _target_: src.models.components.geo_encoders.identity_encoder.IdentityEncoder
+    geo_data_name: aef_avr
+  geo_adapter:
+    _target_: src.models.components.projectors_adapters.mlp_projector.MLPProjector
+    nn_layers: 2
+    hidden_dim: 512
+    #    output_dim: 1280 # TODO for LLM
+    output_dim: 768 # TODO for CLIP text encoder
+  # TODO choose the text encoder
+  text_encoder:
+    _target_: src.models.components.text_encoders.clip_text_encoder.ClipTextEncoder
+    #    _target_: src.models.components.text_encoders.llm2clip_text_encoder.LLM2CLIPTextEncoder
+    hf_cache_dir: ${paths.huggingface_cache}
+    use_geoclip_projector: false # TODO for CLIP encoder
+  # TODO choose the loss fn
+  loss_fn:
+    _target_: src.models.components.loss_fns.soft_contrastive_loss.SoftContrastiveLoss
+    #    _target_: src.models.components.loss_fns.clip_loss.ClipLoss
+    temperature: 0.07
+    stats_file: ${paths.data_dir}/s2bms/aux_stats/unlabelled.json # TODO choose accordingly (for soft loss only)
+    #    stats_file: ${paths.data_dir}/s2bms/aux_stats/labelled.json
+  trainable_modules: [loss_fn.log_temp]
+
+data:
+  dataset:
+    modalities: # TODO
+      aef_avr:
+        path: ${paths.data_dir}/s2bms/eo/avr_aef_128_unlabelled.csv
+        enable_nans: true
+    use_target_data: false
+    use_aux_data: "all"
+    use_unlabelled_data: true # TODO
+    dtype: float32
+  caption_builder:
+    _target_: src.data.butterfly_caption_builder.ButterflyCaptionBuilder
+    templates_fname: v5.json # TODO
+    concepts_fname: v4.json # TODO
+    data_dir: ${paths.data_dir}/s2bms
+    seed: ${seed}
+    n_captions_for_validation: 50
+  saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" # TODO
+  #  saved_split_file_name: "s2bms_union_val_test.pth"
+  num_workers: 14
+  persistent_workers: true
+  pin_memory: true
+  prefetch_factor: 6
+  batch_size: 256 # TODO
+
+logger:
+  wandb:
+    tags: ${tags}
+    group: "alignment"
+    project: "s2bms_alignment"
+  aim:
+    experiment: "alignment"
+
+trainer.num_sanity_val_steps: 0
+
+hydra:
+  mode: MULTIRUN
+  sweeper:
+    params:
+      seed: 12345 #, 404, 654
+      +data.dataset.modalities.aef_avr.size: 128, 256 # For avr (from csv file)
+      #      +data.dataset.modalities.aef.size: 128, 256 # for tiles
+      ++model.loss_fn.sigma: 0.25 #, 0.5, 1
+      ++data.caption_builder.return_aux_ids: true #, false
+      ++callbacks.early_stopping.patience: 20

configs/experiment/s2bms_align_avr_aef_clip_text.yaml

@@ -0,0 +1,82 @@
+# @package _global_
+
+# to execute this experiment run:
+# python train.py experiment=example
+
+# Avr aef 128 + geoclip
+
+defaults:
+  #  - /callbacks@callbacks.batch_size_finder: batch_size_finder
+  - override /model: s2bms_alignment
+  - override /data: s2bms_prediction
+  - override /metrics: s2bms_alignment
+
+tags:
+  - alignment
+  - aef
+  - aef_128
+  - ${str:${seed}}
+
+experiment_name: lab_avr_aef_128_mlp_cliptext
+projects: s2bms_alignment
+
+model:
+  geo_encoder:
+    _target_: src.models.components.geo_encoders.identity_encoder.IdentityEncoder
+    geo_data_name: aef_avr
+  geo_adapter:
+    _target_: src.models.components.projectors_adapters.mlp_projector.MLPProjector
+    nn_layers: 2
+    hidden_dim: 512
+    output_dim: 768
+  text_encoder:
+    _target_: src.models.components.text_encoders.clip_text_encoder.ClipTextEncoder
+    hf_cache_dir: ${paths.huggingface_cache}
+    use_geoclip_projector: false
+  loss_fn:
+    _target_: src.models.components.loss_fns.soft_contrastive_loss.SoftContrastiveLoss
+    temperature: 0.07
+  trainable_modules: [loss_fn.log_temp]
+
+data:
+  dataset:
+    modalities:
+      aef_avr:
+        path: ${paths.data_dir}/s2bms/eo/avr_aef_128.csv
+        enable_nans: true
+    use_target_data: false
+    use_aux_data: "all"
+    dtype: float32
+  caption_builder:
+    _target_: src.data.butterfly_caption_builder.ButterflyCaptionBuilder
+    templates_fname: v5.json
+    concepts_fname: v4.json
+    data_dir: ${paths.data_dir}/s2bms
+    seed: ${seed}
+    n_captions_for_validation: 50
+    stats_file: ${paths.data_dir}/s2bms/aux_stats/unlabelled.json
+
+  num_workers: 14
+  persistent_workers: true
+  pin_memory: true
+  prefetch_factor: 6
+  batch_size: 256
+  saved_split_file_name: "s2bms_aef_coords.pth"
+#    data.saved_split_file_name: "split_indices_s2bms+s2bms-unlabelled-20260602_2026-06-02-0918.pth"
+
+logger:
+  wandb:
+    tags: ${tags}
+    group: "alignment"
+    project: "s2bms_alignment"
+  aim:
+    experiment: "alignment"
+
+hydra:
+  mode: MULTIRUN
+  sweeper:
+    params:
+      seed: 12345 #, 404, 654
+      ++model.loss_fn.sigma: 0.25, 0.5, 1
+      ++data.caption_builder.return_aux_ids: true, false
+      ++callbacks.early_stopping.patience: 20

configs/experiment/s2bms_align_avr_aef_clip_text_1.yaml

@@ -0,0 +1,81 @@
+# @package _global_
+
+# to execute this experiment run:
+# python train.py experiment=example
+
+# Avr aef 128 + geoclip
+
+defaults:
+  #  - /callbacks@callbacks.batch_size_finder: batch_size_finder
+  - override /model: s2bms_alignment
+  - override /data: s2bms_prediction
+  - override /metrics: s2bms_alignment
+
+tags:
+  - alignment
+  - aef
+  - aef_128
+  - ${str:${seed}}
+
+experiment_name: lab_avr_aef_128_mlp_cliptext
+projects: s2bms_alignment
+
+model:
+  geo_encoder:
+    _target_: src.models.components.geo_encoders.identity_encoder.IdentityEncoder
+    geo_data_name: aef_avr
+  geo_adapter:
+    _target_: src.models.components.projectors_adapters.mlp_projector.MLPProjector
+    nn_layers: 2
+    hidden_dim: 512
+    output_dim: 768
+  text_encoder:
+    _target_: src.models.components.text_encoders.clip_text_encoder.ClipTextEncoder
+    hf_cache_dir: ${paths.huggingface_cache}
+    use_geoclip_projector: false
+  loss_fn:
+    _target_: src.models.components.loss_fns.clip_loss.ClipLoss
+    temperature: 0.07
+  trainable_modules: [loss_fn.log_temp]
+
+data:
+  dataset:
+    modalities:
+      aef_avr:
+        path: ${paths.data_dir}/s2bms/eo/avr_aef_128.csv
+        enable_nans: true
+    use_target_data: false
+    use_aux_data: "all"
+    dtype: float32
+  caption_builder:
+    _target_: src.data.butterfly_caption_builder.ButterflyCaptionBuilder
+    templates_fname: v5.json
+    concepts_fname: v4.json
+    data_dir: ${paths.data_dir}/s2bms
+    seed: ${seed}
+    n_captions_for_validation: 50
+    stats_file: ${paths.data_dir}/s2bms/aux_stats/unlabelled.json
+
+  num_workers: 14
+  persistent_workers: true
+  pin_memory: true
+  prefetch_factor: 6
+  batch_size: 256
+  saved_split_file_name: "s2bms_aef_coords.pth"
+#    data.saved_split_file_name: "split_indices_s2bms+s2bms-unlabelled-20260602_2026-06-02-0918.pth"
+
+logger:
+  wandb:
+    tags: ${tags}
+    group: "alignment"
+    project: "s2bms_alignment"
+  aim:
+    experiment: "alignment"
+
+hydra:
+  mode: MULTIRUN
+  sweeper:
+    params:
+      seed: 12345 #, 404, 654
+      ++data.caption_builder.return_aux_ids: false
+      ++callbacks.early_stopping.patience: 20