diff --git a/configs/callbacks/batch_size_finder.yaml b/configs/callbacks/batch_size_finder.yaml index 55533d3..d9c117b 100644 --- a/configs/callbacks/batch_size_finder.yaml +++ b/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 \ No newline at end of file +mode: "power" +init_val: 16 +max_val: 1024 diff --git a/configs/callbacks/early_stopping.yaml b/configs/callbacks/early_stopping.yaml index c826c8d..0ab2df1 100644 --- a/configs/callbacks/early_stopping.yaml +++ b/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 diff --git a/configs/data/butterfly_coords.yaml b/configs/data/butterfly_coords.yaml index f4acb51..f0ad2d4 100644 --- a/configs/data/butterfly_coords.yaml +++ b/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} diff --git a/configs/data/butterfly_coords_text.yaml b/configs/data/butterfly_coords_text.yaml index 52cae6d..5fad06f 100644 --- a/configs/data/butterfly_coords_text.yaml +++ b/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} diff --git a/configs/data/butterfly_coords_text_unlabelled.yaml b/configs/data/butterfly_coords_text_unlabelled.yaml index 9ac97d0..e8399d3 100644 --- a/configs/data/butterfly_coords_text_unlabelled.yaml +++ b/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} diff --git a/configs/data/butterfly_full_param_example.yaml b/configs/data/butterfly_full_param_example.yaml index 78923f5..63320dd 100644 --- a/configs/data/butterfly_full_param_example.yaml +++ b/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} diff --git a/configs/data/butterfly_s2_rgb.yaml b/configs/data/butterfly_s2_rgb.yaml index 1b4a627..4916c57 100644 --- a/configs/data/butterfly_s2_rgb.yaml +++ b/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} diff --git a/configs/experiment/alignment.yaml b/configs/experiment/alignment.yaml index d6281f8..4261a35 100644 --- a/configs/experiment/alignment.yaml +++ b/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} diff --git a/configs/experiment/s2bms_align_avr_aef_clip_text.yaml b/configs/experiment/s2bms_align_avr_aef_clip_text.yaml new file mode 100644 index 0000000..f60723b --- /dev/null +++ b/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 diff --git a/configs/experiment/s2bms_align_avr_aef_clip_text_1.yaml b/configs/experiment/s2bms_align_avr_aef_clip_text_1.yaml new file mode 100644 index 0000000..d5417f8 --- /dev/null +++ b/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 diff --git a/configs/experiment/s2bms_align_avr_aef_clip_text_unlab.yaml b/configs/experiment/s2bms_align_avr_aef_clip_text_unlab.yaml new file mode 100644 index 0000000..f1d13dd --- /dev/null +++ b/configs/experiment/s2bms_align_avr_aef_clip_text_unlab.yaml @@ -0,0 +1,83 @@ +# @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: unlab_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_256_unlabelled.csv + enable_nans: true + use_target_data: false + use_aux_data: "all" + use_unlabelled_data: true + dtype: float32 + return_name_loc: true + 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/labelled.json + + num_workers: 14 + persistent_workers: true + pin_memory: true + prefetch_factor: 6 + batch_size: 256 + saved_split_file_name: "s2bms_10k-unlabelled_coords_aef256.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 diff --git a/configs/experiment/s2bms_align_avr_aef_clip_text_unlab_1.yaml b/configs/experiment/s2bms_align_avr_aef_clip_text_unlab_1.yaml new file mode 100644 index 0000000..892897a --- /dev/null +++ b/configs/experiment/s2bms_align_avr_aef_clip_text_unlab_1.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: unlab_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_256_unlabelled.csv + enable_nans: true + use_target_data: false + use_aux_data: "all" + use_unlabelled_data: true + dtype: float32 + return_name_loc: true + 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/labelled.json + + num_workers: 14 + persistent_workers: true + pin_memory: true + prefetch_factor: 6 + batch_size: 256 + saved_split_file_name: "s2bms_10k-unlabelled_coords_aef256.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 diff --git a/configs/experiment/s2bms_align_avr_aef_llm.yaml b/configs/experiment/s2bms_align_avr_aef_llm.yaml new file mode 100644 index 0000000..e2674da --- /dev/null +++ b/configs/experiment/s2bms_align_avr_aef_llm.yaml @@ -0,0 +1,83 @@ +# @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_llm +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 + text_encoder: + _target_: src.models.components.text_encoders.llm2clip_text_encoder.LLM2CLIPTextEncoder + hf_cache_dir: ${paths.huggingface_cache} + 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/labelled.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" + +trainer.num_sanity_val_steps: 0 + +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 diff --git a/configs/experiment/s2bms_experiment_1.yaml b/configs/experiment/s2bms_experiment_1.yaml index 591b889..6c24802 100644 --- a/configs/experiment/s2bms_experiment_1.yaml +++ b/configs/experiment/s2bms_experiment_1.yaml @@ -38,7 +38,7 @@ data: pin_memory: true num_workers: 14 batch_size: 128 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_10.yaml b/configs/experiment/s2bms_experiment_10.yaml index 09ed4bf..3a69bfd 100644 --- a/configs/experiment/s2bms_experiment_10.yaml +++ b/configs/experiment/s2bms_experiment_10.yaml @@ -43,7 +43,7 @@ data: batch_size: 256 persistent_workers: true prefetch_factor: 2 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_10_1.yaml b/configs/experiment/s2bms_experiment_10_1.yaml index adbef0f..2ae5339 100644 --- a/configs/experiment/s2bms_experiment_10_1.yaml +++ b/configs/experiment/s2bms_experiment_10_1.yaml @@ -43,7 +43,7 @@ data: batch_size: 256 persistent_workers: true prefetch_factor: 2 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_11.yaml b/configs/experiment/s2bms_experiment_11.yaml index 4f0f986..550f2f3 100644 --- a/configs/experiment/s2bms_experiment_11.yaml +++ b/configs/experiment/s2bms_experiment_11.yaml @@ -41,7 +41,7 @@ data: batch_size: 128 persistent_workers: true prefetch_factor: 2 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_12.yaml b/configs/experiment/s2bms_experiment_12.yaml index eb69c23..b208c7e 100644 --- a/configs/experiment/s2bms_experiment_12.yaml +++ b/configs/experiment/s2bms_experiment_12.yaml @@ -41,7 +41,7 @@ data: batch_size: 128 persistent_workers: true prefetch_factor: 2 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test_train.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_13.yaml b/configs/experiment/s2bms_experiment_13.yaml index c71624b..a6872a8 100644 --- a/configs/experiment/s2bms_experiment_13.yaml +++ b/configs/experiment/s2bms_experiment_13.yaml @@ -43,7 +43,7 @@ data: batch_size: 128 persistent_workers: true prefetch_factor: 2 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_14.yaml b/configs/experiment/s2bms_experiment_14.yaml index e5412c7..2de8821 100644 --- a/configs/experiment/s2bms_experiment_14.yaml +++ b/configs/experiment/s2bms_experiment_14.yaml @@ -43,7 +43,7 @@ data: batch_size: 128 persistent_workers: true prefetch_factor: 2 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_15.yaml b/configs/experiment/s2bms_experiment_15.yaml new file mode 100644 index 0000000..0be763a --- /dev/null +++ b/configs/experiment/s2bms_experiment_15.yaml @@ -0,0 +1,71 @@ +# @package _global_ + +# to execute this experiment run: +# python train.py experiment=example + +# Better toegther: AEF with average encoder and satclip, and MLP prediction head, for tile sizes 128, each with 3 random seeds + +defaults: + - override /model: s2bms_prediction + - override /data: s2bms_prediction + - override /metrics: s2bms_predictive + +tags: + - prediction + - aef_satclip + - aef_128 + - ${str:${seed}} + +experiment_name: avr_aef_128_and_satclip_mlp + +model: + geo_encoder: + _target_: src.models.components.geo_encoders.encoder_wrapper.EncoderWrapper + + encoder_branches: + - encoder: + _target_: src.models.components.geo_encoders.identity_encoder.IdentityEncoder + geo_data_name: aef_avr + - encoder: + _target_: src.models.components.geo_encoders.satclip.SatClipCoordinateEncoder + geo_data_name: coords + hf_cache_dir: ${paths.huggingface_cache} + # return_dtype: float32 + normalise: false + fusion_strategy: "concat" + + trainable_modules: [prediction_head] + prediction_head: + _target_: src.models.components.pred_heads.mlp_pred_head.MLPPredictionHead + hidden_dim: 512 + +data: + dataset: + modalities: + coords: + aef_avr: + path: ${paths.data_dir}/s2bms/eo/avr_aef_128.csv + enable_nans: true + use_target_data: true + use_aux_data: false + + persistent_workers: true + pin_memory: true + prefetch_factor: 4 + num_workers: 7 + batch_size: 128 + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" + +logger: + wandb: + tags: ${tags} + group: "predictive" + project: "s2bms_prediction" + aim: + experiment: "predictive" + +hydra: + mode: MULTIRUN + sweeper: + params: + seed: 12345, 404, 654 diff --git a/configs/experiment/s2bms_experiment_2.yaml b/configs/experiment/s2bms_experiment_2.yaml index 0d9a84d..0a1eea9 100644 --- a/configs/experiment/s2bms_experiment_2.yaml +++ b/configs/experiment/s2bms_experiment_2.yaml @@ -38,7 +38,7 @@ data: pin_memory: true num_workers: 14 batch_size: 128 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_3.yaml b/configs/experiment/s2bms_experiment_3.yaml index e4d789f..7f8ce34 100644 --- a/configs/experiment/s2bms_experiment_3.yaml +++ b/configs/experiment/s2bms_experiment_3.yaml @@ -39,7 +39,7 @@ data: num_workers: 14 batch_size: 128 persistent_workers: false - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test_train.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_4.yaml b/configs/experiment/s2bms_experiment_4.yaml index 68dd9f4..3d6d510 100644 --- a/configs/experiment/s2bms_experiment_4.yaml +++ b/configs/experiment/s2bms_experiment_4.yaml @@ -25,6 +25,7 @@ model: trainable_modules: [prediction_head] prediction_head: _target_: src.models.components.pred_heads.linear_pred_head.LinearPredictionHead + hidden_dim: 512 data: dataset: @@ -39,7 +40,7 @@ data: num_workers: 14 batch_size: 128 persistent_workers: false - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test_train.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_5.yaml b/configs/experiment/s2bms_experiment_5.yaml index 366bfed..2b6f2c2 100644 --- a/configs/experiment/s2bms_experiment_5.yaml +++ b/configs/experiment/s2bms_experiment_5.yaml @@ -40,7 +40,7 @@ data: num_workers: 14 batch_size: 128 persistent_workers: false - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test_train.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_6.yaml b/configs/experiment/s2bms_experiment_6.yaml index 2cd63c8..66d0cf0 100644 --- a/configs/experiment/s2bms_experiment_6.yaml +++ b/configs/experiment/s2bms_experiment_6.yaml @@ -34,7 +34,7 @@ data: caption_builder: num_workers: 14 batch_size: 128 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_7.yaml b/configs/experiment/s2bms_experiment_7.yaml index ec4f230..c0ec25f 100644 --- a/configs/experiment/s2bms_experiment_7.yaml +++ b/configs/experiment/s2bms_experiment_7.yaml @@ -34,7 +34,7 @@ data: caption_builder: num_workers: 14 batch_size: 128 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_8.yaml b/configs/experiment/s2bms_experiment_8.yaml index 4af78f0..50f9d63 100644 --- a/configs/experiment/s2bms_experiment_8.yaml +++ b/configs/experiment/s2bms_experiment_8.yaml @@ -36,7 +36,7 @@ data: caption_builder: num_workers: 14 batch_size: 128 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/experiment/s2bms_experiment_9.yaml b/configs/experiment/s2bms_experiment_9.yaml index 4bd5def..d5bb66d 100644 --- a/configs/experiment/s2bms_experiment_9.yaml +++ b/configs/experiment/s2bms_experiment_9.yaml @@ -36,7 +36,7 @@ data: caption_builder: num_workers: 14 batch_size: 128 - saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" + saved_split_file_name: "s2bms_unlabelled_union_val_test.pth" logger: wandb: diff --git a/configs/inference.yaml b/configs/inference.yaml index 45f5bcc..47ad5e4 100644 --- a/configs/inference.yaml +++ b/configs/inference.yaml @@ -1,12 +1,14 @@ defaults: - - _self_ - paths: ${oc.env:STORAGE_MODE,local} - extras: default - hydra: default + - data: s2bms_prediction + - _self_ # task name, determines output directory path task_name: "inference" tags: ["dev"] +seed: 12345 # If set, inference.py will load this merged checkpoint directly. #inference_ckpt_path: ${paths.log_dir}${task_name}/2026-04-07_12-56-16/inference_model.ckpt @@ -14,12 +16,37 @@ tags: ["dev"] # If `inference_ckpt_path` is not set, stitch the inference model from: # - predictive ckpt (provides prediction_head weights) # - alignment ckpt (provides text_encoder weights + geo_encoder) -predictive_ckpt_path: ${paths.log_dir}train/runs/2026-04-02_15-54-53/checkpoints/epoch_000.ckpt -alignment_ckpt_path: ${paths.log_dir}train/runs/2026-04-02_15-40-03/checkpoints/epoch_000.ckpt +predictive_ckpt_path: ${paths.checkpoint_root}/s2bms_prediction_ckpt/1cti8qah_epoch_037.ckpt +alignment_ckpt_path: ${paths.checkpoint_root}/s2bms_alignment_ckpt/dmts2wpm_epoch_097.ckpt # If set, inference.py will save a merged inference checkpoint you can reload with # `inference_ckpt_path`. -save_inference_ckpt_path: ${paths.log_dir}${task_name}/${now:%Y-%m-%d}_${now:%H-%M-%S}/inference_model.ckpt +#save_inference_ckpt_path: ${paths.checkpoint_root}/s2bms_inference_model/${now:%Y-%m-%d}_${now:%H-%M-%S}.ckpt +inference_ckpt_path: ${paths.checkpoint_root}/s2bms_inference_model/2026-06-09_13-17-38.ckpt #save_inference_ckpt_path: null +save_ckpt: true + +training_order: ["prediction_model", "alignment_model"] + +data: + dataset: + modalities: + aef_avr: + path: ${paths.data_dir}/s2bms/eo/avr_aef_128.csv + enable_nans: true -training_order: ["alignment_model", "prediction_model"] + use_target_data: false + use_aux_data: "all" + dtype: float32 + return_name_loc: true + 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} + num_workers: 8 + persistent_workers: false + pin_memory: false + batch_size: 64 + saved_split_file_name: "split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth" diff --git a/configs/paths/local.yaml b/configs/paths/local.yaml index c3683f3..c7e598e 100644 --- a/configs/paths/local.yaml +++ b/configs/paths/local.yaml @@ -23,4 +23,4 @@ work_dir: ${hydra:runtime.cwd} # huggingface cache directory # can be overridden via HF_HOME environment variable -huggingface_cache: ${oc.env:HF_HOME,${paths.root_dir}/.cache/huggingface} +huggingface_cache: ${oc.env:HF_HOME,${paths.cache_dir}/.cache/huggingface} diff --git a/configs/train.yaml b/configs/train.yaml index 894cfa0..5faeeca 100644 --- a/configs/train.yaml +++ b/configs/train.yaml @@ -43,6 +43,8 @@ train: True # lightning chooses best weights based on the metric specified in checkpoint callback test: True +validate: True + # simply provide checkpoint path to resume training ckpt_path: null diff --git a/data/s2bms/aux_stats/explanations b/data/s2bms/aux_stats/explanations new file mode 100644 index 0000000..53dc9f6 --- /dev/null +++ b/data/s2bms/aux_stats/explanations @@ -0,0 +1,2 @@ +labelled - s2bms aux statistical values, train split (from model_ready csv used) +unlabelled - model_ready_s2bms-unlabelled-20260529.csv aux col stats for training data split. diff --git a/data/s2bms/aux_stats/labelled.json b/data/s2bms/aux_stats/labelled.json new file mode 100644 index 0000000..60a2845 --- /dev/null +++ b/data/s2bms/aux_stats/labelled.json @@ -0,0 +1,350 @@ +{ + "aux_bioclim_01": { + "mean": 9.478358208955223, + "std": 0.6813968574277965 + }, + "aux_bioclim_02": { + "mean": 7.261727078891258, + "std": 0.5743689424967182 + }, + "aux_bioclim_03": { + "mean": 35.45842217484009, + "std": 1.5745028402780046 + }, + "aux_bioclim_04": { + "mean": 46.482974413646055, + "std": 2.7476746356316304 + }, + "aux_bioclim_05": { + "mean": 20.681236673773988, + "std": 1.1951117823510142 + }, + "aux_bioclim_06": { + "mean": 0.4868869936034116, + "std": 0.7787387594816507 + }, + "aux_bioclim_07": { + "mean": 20.19434968017058, + "std": 1.3614215267253051 + }, + "aux_bioclim_08": { + "mean": 6.26321961620469, + "std": 2.734385392685174 + }, + "aux_bioclim_09": { + "mean": 9.148614072494668, + "std": 3.74705621891078 + }, + "aux_bioclim_10": { + "mean": 15.553304904051172, + "std": 0.8357249773564381 + }, + "aux_bioclim_11": { + "mean": 3.7204690831556504, + "std": 0.7173613109007019 + }, + "aux_bioclim_12": { + "mean": 775.1769722814499, + "std": 168.33111593490483 + }, + "aux_bioclim_13": { + "mean": 82.63432835820896, + "std": 22.906370014077783 + }, + "aux_bioclim_14": { + "mean": 47.988272921108745, + "std": 9.389409722851463 + }, + "aux_bioclim_15": { + "mean": 16.553304904051174, + "std": 5.307596573012521 + }, + "aux_bioclim_16": { + "mean": 237.5095948827292, + "std": 65.82287942145929 + }, + "aux_bioclim_17": { + "mean": 156.3582089552239, + "std": 26.224456836761853 + }, + "aux_bioclim_18": { + "mean": 175.53411513859274, + "std": 29.036381115737576 + }, + "aux_bioclim_19": { + "mean": 211.044776119403, + "std": 62.222859349413106 + }, + "aux_corine_frac_1": { + "mean": 0.1974766383028437, + "std": 0.2602392974374031 + }, + "aux_corine_frac_11": { + "mean": 0.1301962252180134, + "std": 0.19606981059582396 + }, + "aux_corine_frac_111": { + "mean": 0.0018835987151241686, + "std": 0.014460454357011451 + }, + "aux_corine_frac_112": { + "mean": 0.12831262650288924, + "std": 0.19143236436257338 + }, + "aux_corine_frac_12": { + "mean": 0.02062538765029979, + "std": 0.05555662489196779 + }, + "aux_corine_frac_121": { + "mean": 0.01621736710326999, + "std": 0.045436911401837476 + }, + "aux_corine_frac_122": { + "mean": 0.0008087694633551867, + "std": 0.0073899002568070336 + }, + "aux_corine_frac_123": { + "mean": 0.0007881621088882273, + "std": 0.00948945250449678 + }, + "aux_corine_frac_124": { + "mean": 0.002811088974786385, + "std": 0.026763607274160314 + }, + "aux_corine_frac_13": { + "mean": 0.004779012207277344, + "std": 0.020532901676720994 + }, + "aux_corine_frac_131": { + "mean": 0.003592628135133028, + "std": 0.018691388297068547 + }, + "aux_corine_frac_132": { + "mean": 0.0005700646265308338, + "std": 0.006091396436645145 + }, + "aux_corine_frac_133": { + "mean": 0.0006163194456134819, + "std": 0.005354453536589029 + }, + "aux_corine_frac_14": { + "mean": 0.04187601322725316, + "std": 0.09053515666961054 + }, + "aux_corine_frac_141": { + "mean": 0.008769670818258761, + "std": 0.039855301035569904 + }, + "aux_corine_frac_142": { + "mean": 0.03310634240899439, + "std": 0.0774796864288167 + }, + "aux_corine_frac_2": { + "mean": 0.5492999788267231, + "std": 0.2885341556305182 + }, + "aux_corine_frac_21": { + "mean": 0.2751627934613993, + "std": 0.26029991741188563 + }, + "aux_corine_frac_211": { + "mean": 0.2751627934613993, + "std": 0.26029991741188563 + }, + "aux_corine_frac_212": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_213": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_22": { + "mean": 0.00041510124495720943, + "std": 0.006567155612433827 + }, + "aux_corine_frac_221": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_222": { + "mean": 0.00041510124495720943, + "std": 0.006567155612433827 + }, + "aux_corine_frac_223": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_23": { + "mean": 0.2624442076480314, + "std": 0.21558169533315324 + }, + "aux_corine_frac_231": { + "mean": 0.2624442076480314, + "std": 0.21558169533315324 + }, + "aux_corine_frac_24": { + "mean": 0.011277876472335122, + "std": 0.05090940063610321 + }, + "aux_corine_frac_241": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_242": { + "mean": 0.0007546994246533241, + "std": 0.010810046743945034 + }, + "aux_corine_frac_243": { + "mean": 0.010523177047681796, + "std": 0.04990143178437093 + }, + "aux_corine_frac_244": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_3": { + "mean": 0.19202690040482187, + "std": 0.22118636236252417 + }, + "aux_corine_frac_31": { + "mean": 0.13380397473440228, + "std": 0.17526935775189442 + }, + "aux_corine_frac_311": { + "mean": 0.07779569292193297, + "std": 0.11887832359511401 + }, + "aux_corine_frac_312": { + "mean": 0.024288529798213176, + "std": 0.07745327927661465 + }, + "aux_corine_frac_313": { + "mean": 0.031719752014256114, + "std": 0.08775508006760226 + }, + "aux_corine_frac_32": { + "mean": 0.057084137390655605, + "std": 0.13510941488940179 + }, + "aux_corine_frac_321": { + "mean": 0.022382606501706252, + "std": 0.08459439268533137 + }, + "aux_corine_frac_322": { + "mean": 0.030206032609034834, + "std": 0.09612787375800697 + }, + "aux_corine_frac_323": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_324": { + "mean": 0.0044954982799145105, + "std": 0.022499180627468546 + }, + "aux_corine_frac_33": { + "mean": 0.0011387882797639548, + "std": 0.011249341135396742 + }, + "aux_corine_frac_331": { + "mean": 0.00020365428949126714, + "std": 0.003671268959218654 + }, + "aux_corine_frac_332": { + "mean": 0.00021130029530543496, + "std": 0.0036542400691276906 + }, + "aux_corine_frac_333": { + "mean": 0.0007238336949672526, + "std": 0.010020154855462676 + }, + "aux_corine_frac_334": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_335": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_4": { + "mean": 0.023167516477294505, + "std": 0.08966431628095389 + }, + "aux_corine_frac_41": { + "mean": 0.010217631418879294, + "std": 0.05956890217531605 + }, + "aux_corine_frac_411": { + "mean": 0.005128956103351367, + "std": 0.03182072782777159 + }, + "aux_corine_frac_412": { + "mean": 0.0050886753155279275, + "std": 0.05078978822825269 + }, + "aux_corine_frac_42": { + "mean": 0.012949885058415208, + "std": 0.06281313316391562 + }, + "aux_corine_frac_421": { + "mean": 0.003738002199751277, + "std": 0.027210848090972052 + }, + "aux_corine_frac_422": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_423": { + "mean": 0.009211882858663931, + "std": 0.04829498490397355 + }, + "aux_corine_frac_5": { + "mean": 0.0380289659883168, + "std": 0.11486072053644752 + }, + "aux_corine_frac_51": { + "mean": 0.0077577450917248325, + "std": 0.03683904294786193 + }, + "aux_corine_frac_511": { + "mean": 0.0006007572538380269, + "std": 0.008124456212658684 + }, + "aux_corine_frac_512": { + "mean": 0.007156987837886806, + "std": 0.036051583858167306 + }, + "aux_corine_frac_52": { + "mean": 0.030271220896591972, + "std": 0.11075559157067964 + }, + "aux_corine_frac_521": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_522": { + "mean": 0.0034053245323920935, + "std": 0.030152244642461808 + }, + "aux_corine_frac_523": { + "mean": 0.02686589636419988, + "std": 0.10564500727708297 + }, + "aux_maxdist_road": { + "mean": 1470.5671641791046, + "std": 746.0098613291074 + }, + "aux_meandist_road": { + "mean": 571.22921108742, + "std": 534.4224060908446 + }, + "aux_pop_density": { + "mean": 543.4296375266524, + "std": 1174.8791303921573 + }, + "aux_total_population": { + "mean": 3544.597014925373, + "std": 7657.245673883998 + } +} diff --git a/data/s2bms/aux_stats/unlabelled.json b/data/s2bms/aux_stats/unlabelled.json new file mode 100644 index 0000000..236fa13 --- /dev/null +++ b/data/s2bms/aux_stats/unlabelled.json @@ -0,0 +1,350 @@ +{ + "aux_bioclim_01": { + "mean": 8.39215464935798, + "std": 1.3489076187162317 + }, + "aux_bioclim_02": { + "mean": 6.822435445181319, + "std": 0.630031288756174 + }, + "aux_bioclim_03": { + "mean": 35.273035134753776, + "std": 1.9165734744337644 + }, + "aux_bioclim_04": { + "mean": 44.353753351206436, + "std": 3.9291740967815807 + }, + "aux_bioclim_05": { + "mean": 18.885409905460705, + "std": 1.8751171337662043 + }, + "aux_bioclim_06": { + "mean": -0.20960914350218712, + "std": 1.3725400398778722 + }, + "aux_bioclim_07": { + "mean": 19.095019048962893, + "std": 1.7191272659843762 + }, + "aux_bioclim_08": { + "mean": 5.64862424156907, + "std": 3.456703372551364 + }, + "aux_bioclim_09": { + "mean": 8.581755326654438, + "std": 2.7864265921108977 + }, + "aux_bioclim_10": { + "mean": 14.232030478340622, + "std": 1.479657387312894 + }, + "aux_bioclim_11": { + "mean": 2.9279243685621563, + "std": 1.3246359592065378 + }, + "aux_bioclim_12": { + "mean": 986.9657118667984, + "std": 342.6762922968029 + }, + "aux_bioclim_13": { + "mean": 108.88739946380697, + "std": 46.04613668954425 + }, + "aux_bioclim_14": { + "mean": 58.09030619444052, + "std": 16.423198571455718 + }, + "aux_bioclim_15": { + "mean": 18.988429518837307, + "std": 5.836359255751771 + }, + "aux_bioclim_16": { + "mean": 311.777762099619, + "std": 128.6739578801313 + }, + "aux_bioclim_17": { + "mean": 185.6207139833498, + "std": 48.97990233262024 + }, + "aux_bioclim_18": { + "mean": 216.47199096938056, + "std": 54.45975422898328 + }, + "aux_bioclim_19": { + "mean": 277.92733173416116, + "std": 119.56237681535922 + }, + "aux_corine_frac_111": { + "mean": 0.0013602792796026007, + "std": 0.019313170626270262 + }, + "aux_corine_frac_112": { + "mean": 0.050006013167217804, + "std": 0.13148336944272226 + }, + "aux_corine_frac_121": { + "mean": 0.009250862708691302, + "std": 0.03881985995144867 + }, + "aux_corine_frac_122": { + "mean": 0.0007880240503307102, + "std": 0.00778518077993335 + }, + "aux_corine_frac_123": { + "mean": 0.00046706624328041413, + "std": 0.009420263901164894 + }, + "aux_corine_frac_124": { + "mean": 0.0011089385169889571, + "std": 0.014364655530256759 + }, + "aux_corine_frac_131": { + "mean": 0.002571445974733077, + "std": 0.02003256510611235 + }, + "aux_corine_frac_132": { + "mean": 0.00017665416030145643, + "std": 0.0035253961655689545 + }, + "aux_corine_frac_133": { + "mean": 0.0005756500688074249, + "std": 0.007402470282847353 + }, + "aux_corine_frac_141": { + "mean": 0.0028120972814816856, + "std": 0.021223127264056292 + }, + "aux_corine_frac_142": { + "mean": 0.011166278833193702, + "std": 0.038165164275206914 + }, + "aux_corine_frac_211": { + "mean": 0.2586450775682287, + "std": 0.32966642776729954 + }, + "aux_corine_frac_212": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_213": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_221": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_222": { + "mean": 0.0003398871933901921, + "std": 0.008045773480885858 + }, + "aux_corine_frac_223": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_231": { + "mean": 0.27208956545103624, + "std": 0.29460130010474606 + }, + "aux_corine_frac_241": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_242": { + "mean": 0.0015301132476283517, + "std": 0.022541814707966536 + }, + "aux_corine_frac_243": { + "mean": 0.0053086880969759696, + "std": 0.029445598088384678 + }, + "aux_corine_frac_244": { + "mean": 1.4017089279405644e-05, + "std": 0.0011800191515226413 + }, + "aux_corine_frac_311": { + "mean": 0.020470399162633234, + "std": 0.052253302577792045 + }, + "aux_corine_frac_312": { + "mean": 0.049353532759643035, + "std": 0.12652637234654923 + }, + "aux_corine_frac_313": { + "mean": 0.01244167875010827, + "std": 0.046049545976652805 + }, + "aux_corine_frac_321": { + "mean": 0.059258514947408365, + "std": 0.14397917118825002 + }, + "aux_corine_frac_322": { + "mean": 0.07682161065734752, + "std": 0.17283380584540975 + }, + "aux_corine_frac_323": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_324": { + "mean": 0.016625096426279753, + "std": 0.05992790760448416 + }, + "aux_corine_frac_331": { + "mean": 0.0004443415806525385, + "std": 0.011574430529737981 + }, + "aux_corine_frac_332": { + "mean": 0.0006745155596987792, + "std": 0.011643913566303933 + }, + "aux_corine_frac_333": { + "mean": 0.011641901867519025, + "std": 0.06223229237384561 + }, + "aux_corine_frac_334": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_335": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_411": { + "mean": 0.0014524852178632711, + "std": 0.020316740577884202 + }, + "aux_corine_frac_412": { + "mean": 0.10249822833445148, + "std": 0.24051158401945827 + }, + "aux_corine_frac_421": { + "mean": 0.001392197302331657, + "std": 0.01753513296020852 + }, + "aux_corine_frac_422": { + "mean": 0.0, + "std": 0.0 + }, + "aux_corine_frac_423": { + "mean": 0.004125761567925976, + "std": 0.03422460913306942 + }, + "aux_corine_frac_511": { + "mean": 0.0005947060170377448, + "std": 0.009810879261322713 + }, + "aux_corine_frac_512": { + "mean": 0.008423078073735008, + "std": 0.05239152569781716 + }, + "aux_corine_frac_521": { + "mean": 2.05628405890116e-05, + "std": 0.0014344231921542916 + }, + "aux_corine_frac_522": { + "mean": 0.002235854221870117, + "std": 0.02562336613443482 + }, + "aux_corine_frac_523": { + "mean": 0.013314875781737049, + "std": 0.06729427140720412 + }, + "aux_corine_frac_1": { + "mean": 0.08028331028462915, + "std": 0.1825027053945205 + }, + "aux_corine_frac_2": { + "mean": 0.5379273486465389, + "std": 0.3899391765334726 + }, + "aux_corine_frac_3": { + "mean": 0.24773159171129058, + "std": 0.31469171199310836 + }, + "aux_corine_frac_4": { + "mean": 0.1094686724225724, + "std": 0.24278980948174828 + }, + "aux_corine_frac_5": { + "mean": 0.024589076934968934, + "std": 0.08840739214247646 + }, + "aux_corine_frac_11": { + "mean": 0.0513662924468204, + "std": 0.13701592738232227 + }, + "aux_corine_frac_12": { + "mean": 0.011614891519291386, + "std": 0.0460258628072768 + }, + "aux_corine_frac_13": { + "mean": 0.003323750203841959, + "std": 0.021796606691765877 + }, + "aux_corine_frac_14": { + "mean": 0.013978376114675386, + "std": 0.0452558663201929 + }, + "aux_corine_frac_21": { + "mean": 0.2586450775682287, + "std": 0.32966642776729954 + }, + "aux_corine_frac_22": { + "mean": 0.0003398871933901921, + "std": 0.008045773480885858 + }, + "aux_corine_frac_23": { + "mean": 0.27208956545103624, + "std": 0.29460130010474606 + }, + "aux_corine_frac_24": { + "mean": 0.0068528184338837285, + "std": 0.03744219094823718 + }, + "aux_corine_frac_31": { + "mean": 0.08226561067238455, + "std": 0.14632419826435697 + }, + "aux_corine_frac_32": { + "mean": 0.1527052220310357, + "std": 0.2502187615368793 + }, + "aux_corine_frac_33": { + "mean": 0.012760759007870343, + "std": 0.06572354806621297 + }, + "aux_corine_frac_41": { + "mean": 0.10395071355231476, + "std": 0.24079581091787888 + }, + "aux_corine_frac_42": { + "mean": 0.005517958870257633, + "std": 0.042508099721360604 + }, + "aux_corine_frac_51": { + "mean": 0.009017784090772753, + "std": 0.05325796209639183 + }, + "aux_corine_frac_52": { + "mean": 0.01557129284419618, + "std": 0.07195454194364073 + }, + "aux_total_population": { + "mean": 1394.7764921687597, + "std": 4643.2852761894 + }, + "aux_pop_density": { + "mean": 213.62537039650064, + "std": 712.4218778428469 + }, + "aux_maxdist_road": { + "mean": 2232.251869620432, + "std": 1319.1743912122638 + }, + "aux_meandist_road": { + "mean": 1239.2566671370114, + "std": 1311.269053382563 + } +} diff --git a/notebooks/09-GT-aef-tessera-datasplit-filtering.ipynb b/notebooks/09-GT-aef-tessera-datasplit-filtering.ipynb index 020dcf6..6c7ffcf 100644 --- a/notebooks/09-GT-aef-tessera-datasplit-filtering.ipynb +++ b/notebooks/09-GT-aef-tessera-datasplit-filtering.ipynb @@ -318,13 +318,7 @@ "source": [] } ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - } - }, + "metadata": {}, "nbformat": 4, "nbformat_minor": 5 } diff --git a/notebooks/10-GT-avr_encoding.ipynb b/notebooks/10-GT-avr_encoding.ipynb new file mode 100644 index 0000000..296c737 --- /dev/null +++ b/notebooks/10-GT-avr_encoding.ipynb @@ -0,0 +1,118 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "import glob\n", + "import os\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "import torch\n", + "\n", + "from src.utils.data_utils import center_crop_npy\n", + "\n", + "os.chdir(\"..\")\n", + "os.getcwd()" + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "# Calculating from tiles" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "mod = \"tessera\"\n", + "size = 128\n", + "\n", + "paths = glob.glob(f\"data/s2bms/eo/{mod}/*UKBMS*.npy\")\n", + "rows = []\n", + "for p in paths:\n", + " arr = np.load(p).transpose(2, 0, 1)\n", + " if arr.dtype != np.dtype(\"float32\"):\n", + " arr = arr.astype(dtype=\"float32\", copy=False)\n", + " arr = center_crop_npy(arr, (64 if mod == \"aef\" else 128, size, size))\n", + " if np.isinf(arr).any():\n", + " arr[np.isinf(arr)] = np.nan\n", + " tensor = torch.from_numpy(arr)\n", + " row = {f\"emb_{i}\": v.item() for i, v in enumerate(tensor.nanmean(dim=(-2, -1)))}\n", + " row[\"name_loc\"] = p.split(\"/\")[-1].split(\".\")[0].replace(f\"{mod}_\", \"\")\n", + " rows.append(row)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "df = pd.DataFrame(rows)\n", + "df" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "if not os.path.exists(f\"data/s2bms/eo/avr_{mod}_{size}.csv\"):\n", + " df.to_csv(f\"data/s2bms/eo/avr_{mod}_{size}.csv\", index=False)" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "# Merging-in unlabelled" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "size = 128\n", + "mod = \"aef\"\n", + "\n", + "df = pd.read_csv(f\"data/s2bms/eo/avr_{mod}_{size}.csv\")\n", + "df_unlabelled = pd.read_csv(f\"data/s2bms/eo/avr_{mod}_{size}_just_unlabelled.csv\")\n", + "\n", + "df_merged = pd.concat([df, df_unlabelled])\n", + "df_merged.name_loc\n", + "\n", + "if not os.path.exists(f\"data/s2bms/eo/avr_{mod}_{size}_unlabelled.csv\"):\n", + " df.to_csv(f\"data/s2bms/eo/avr_{mod}_{size}_unlabelled.csv\")\n", + "else:\n", + " print(\"Already saved\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/11-GT-experimenting-with-text-embeddings.ipynb b/notebooks/11-GT-experimenting-with-text-embeddings.ipynb new file mode 100644 index 0000000..91e2aa8 --- /dev/null +++ b/notebooks/11-GT-experimenting-with-text-embeddings.ipynb @@ -0,0 +1,302 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "from dotenv import load_dotenv\n", + "\n", + "load_dotenv()\n", + "os.chdir(\"..\")\n", + "os.getcwd()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import itertools\n", + "import os\n", + "\n", + "import hydra\n", + "import lightning as pl\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import seaborn as sns\n", + "import torch\n", + "import torch.nn.functional as F\n", + "from lightning import LightningDataModule, Trainer\n", + "\n", + "if os.environ.get(\"TOKENIZERS_PARALLELISM\") is None:\n", + " os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from hydra import compose, initialize\n", + "from hydra.core.hydra_config import HydraConfig\n", + "\n", + "with initialize(version_base=\"1.3\", config_path=\"../configs\"):\n", + " cfg = compose(\n", + " config_name=\"train.yaml\", return_hydra_config=True, overrides=[\"experiment=pycharm\"]\n", + " )\n", + "\n", + "HydraConfig.instance().set_config(cfg)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data)\n", + "datamodule.setup()\n", + "\n", + "trainer: Trainer = pl.Trainer(accelerator=\"mps\", max_epochs=1)\n", + "trainer.datamodule = datamodule\n", + "\n", + "# Get model\n", + "model = hydra.utils.instantiate(cfg.model)\n", + "model.to(\"mps\")\n", + "model.trainer = trainer\n", + "model.setup(\"test\")\n", + "\n", + "is_llm = \"LLM\" in cfg.model.text_encoder._target_" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "captions = [c[\"concept_caption\"] for c in datamodule.caption_builder.__dict__[\"concepts\"]]\n", + "b = {\"text\": captions}\n", + "text_feats = model.text_encoder(b, \"train\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "text_feats = F.normalize(text_feats)\n", + "cos_sim = text_feats @ text_feats.T\n", + "if is_llm:\n", + " cos_sim = cos_sim.detach().cpu().to(torch.float32).numpy()\n", + "else:\n", + " cos_sim = cos_sim.detach().cpu().detach().numpy()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "fig = plt.figure(figsize=(16, 8))\n", + "gs = fig.add_gridspec(1, 2, width_ratios=[1.5, 1])\n", + "\n", + "ax1 = fig.add_subplot(gs[0])\n", + "short_labels = [f\"C{i}\" for i in range(len(captions))]\n", + "\n", + "sns.heatmap(\n", + " cos_sim,\n", + " xticklabels=short_labels,\n", + " yticklabels=short_labels,\n", + " ax=ax1,\n", + " cmap=\"viridis\",\n", + " # vmin=-1, vmax=1,\n", + " cbar_kws={\"label\": \"Cosine Similarity\"},\n", + ")\n", + "\n", + "ax1.set_title(\n", + " f\"Cosine Similarity Matrix For Concepts With {'LLM' if is_llm else 'CLIP'} Text Encoder\"\n", + ")\n", + "\n", + "ax2 = fig.add_subplot(gs[1])\n", + "ax2.axis(\"off\")\n", + "\n", + "legend_text = \"\\n\".join([f\"C{i}: {cap}\" for i, cap in enumerate(captions)])\n", + "ax2.text(0, 1, legend_text, fontsize=10, va=\"top\", ha=\"left\", wrap=True)\n", + "\n", + "plt.tight_layout()\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "dataloader = datamodule.train_dataloader()\n", + "dataloader_cycle = itertools.cycle(dataloader)\n", + "\n", + "all_captions = []\n", + "all_embeddings = []\n", + "total_samples = 0\n", + "max_seq_len = 0\n", + "target_samples = 10000" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "for batch in dataloader_cycle:\n", + " with torch.no_grad():\n", + " model.eval()\n", + " if is_llm:\n", + " text_input = batch.get(\"text\")\n", + " text_feats = model.text_encoder(batch, \"train\")\n", + " else:\n", + " text_input = batch.get(\"text\")\n", + " text_tokens = model.text_encoder.processor(\n", + " text=text_input,\n", + " return_tensors=\"pt\",\n", + " padding=True,\n", + " truncation=True,\n", + " max_length=None,\n", + " )\n", + " if text_tokens.input_ids.shape[-1] <= 77:\n", + " device = \"mps\"\n", + " text_tokens = {k: v.to(device) for k, v in text_tokens.items()}\n", + " text_embeds = model.text_encoder.model.get_text_features(**text_tokens)\n", + " if model.text_encoder.projector is not None:\n", + " text_embeds = model.text_encoder.projector(text_embeds)\n", + " if model.text_encoder.extra_projector is not None:\n", + " text_embeds = model.text_encoder.extra_projector(text_embeds)\n", + " else:\n", + " for i, t in enumerate(text_input):\n", + " print(t, len(text_tokens[\"input_ids\"][i]))\n", + "\n", + " if len(text_tokens[\"input_ids\"][i]) > max_seq_len:\n", + " max_seq_len = len(text_tokens[\"input_ids\"][i])\n", + " break\n", + "\n", + " all_embeddings.append(text_feats.cpu())\n", + " all_captions.extend(batch[\"text\"])\n", + "\n", + " total_samples += text_feats.shape[0]\n", + " print(total_samples)\n", + "\n", + " # Break only when we have enough\n", + " if total_samples >= target_samples:\n", + " break\n", + "\n", + "# 3. Finalize\n", + "text_feats = torch.cat(all_embeddings, dim=0)[:target_samples]\n", + "captions = all_captions[:target_samples]\n", + "\n", + "text_feats = F.normalize(text_feats)\n", + "cos_sim = text_feats @ text_feats.T" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "captions_arr = np.array(captions)\n", + "identity_mask = captions_arr[:, None] == captions_arr[None, :]\n", + "lower_tri_mask = np.tril(np.ones(identity_mask.shape, dtype=bool), k=-1)\n", + "\n", + "final_mask = (~identity_mask) & lower_tri_mask\n", + "if is_llm:\n", + " unique_sims = cos_sim.detach().cpu().to(torch.float32).numpy()[final_mask]\n", + "else:\n", + " unique_sims = cos_sim.detach().cpu().numpy()[final_mask]\n", + "\n", + "row_idx, col_idx = np.where(final_mask)\n", + "\n", + "# Plot\n", + "plt.figure(figsize=(8, 5))\n", + "sns.histplot(unique_sims, bins=50, kde=True)\n", + "plt.title(\n", + " f\"Distribution of Cosine Similarities \\n(for random 10k training location captions with {'LLM' if is_llm else 'CLIP'} text encoder)\"\n", + ")\n", + "plt.xlabel(\"Cosine Similarity\")\n", + "plt.ylabel(\"Frequency\")\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "# Get the lower triangle\n", + "lower_tri_indices = np.tril_indices(cos_sim.shape[0], k=-1)\n", + "\n", + "# Get indices of the unique pairs\n", + "row_idx, col_idx = lower_tri_indices\n", + "\n", + "# Sort the unique similarities\n", + "sorted_indices = np.argsort(unique_sims) # Ascending order\n", + "\n", + "print(\"--- Highest Similarity Pairs ---\")\n", + "for i in sorted_indices[-5:][::-1]:\n", + " r, c = row_idx[i], col_idx[i]\n", + " print(f\"Score: {unique_sims[i]:.4f}\")\n", + " print(f\" - '{captions[r]}'\")\n", + " print(f\" - '{captions[c]}'\\n\")\n", + "\n", + "print(\"--- Lowest Similarity Pairs ---\")\n", + "for i in sorted_indices[:5]:\n", + " r, c = row_idx[i], col_idx[i]\n", + " print(f\"Score: {unique_sims[i]:.4f}\")\n", + " print(f\" - '{captions[r]}'\")\n", + " print(f\" - '{captions[c]}'\\n\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 2 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython2", + "version": "2.7.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/12-GT-aux-stats.ipynb b/notebooks/12-GT-aux-stats.ipynb new file mode 100644 index 0000000..d7bc8cf --- /dev/null +++ b/notebooks/12-GT-aux-stats.ipynb @@ -0,0 +1,124 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import os\n", + "\n", + "import pandas as pd\n", + "import torch\n", + "\n", + "os.chdir(\"..\")\n", + "os.getcwd()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "labelled = True\n", + "\n", + "output_dir = \"/Users/gabriele/Desktop/aether/data/s2bms/aux_stats/\"\n", + "os.makedirs(output_dir, exist_ok=True)\n", + "\n", + "if labelled:\n", + " df = pd.read_csv(\"/Users/gabriele/Desktop/aether/data/s2bms/model_ready_s2bms.csv\")\n", + " output_path = output_dir + \"labelled.json\"\n", + "else:\n", + " df = pd.read_csv(\n", + " \"/Users/gabriele/Desktop/aether/data/s2bms/model_ready_s2bms-unlabelled-20260529.csv\"\n", + " )\n", + " output_path = output_dir + \"unlabelled.json\"\n", + "\n", + "col = [i for i in df.columns if \"aux\" in i and \"top\" not in i]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "# Load in original datasplit file\n", + "if labelled:\n", + " pth = \"data/s2bms/splits/s2bms_union_val_test.pth\"\n", + "else:\n", + " pth = \"data/s2bms/splits/s2bms_unlabelled_union_val_test.pth\"\n", + "\n", + "\n", + "split_indices = torch.load(pth, weights_only=False)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "df = df[df[\"name_loc\"].isin(split_indices[\"train_indices\"])]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "if os.path.exists(output_path):\n", + " print(f\"Output file already exists: {output_path}\")\n", + "else:\n", + " stats = {}\n", + " for col in col:\n", + " stats[col] = {\n", + " \"mean\": float(df[col].mean()),\n", + " \"std\": float(df[col].std()),\n", + " }\n", + "\n", + " with open(output_path, \"w\") as f:\n", + " json.dump(stats, f, indent=4)\n", + "\n", + " print(f\"Saved global statistics for {len(stats)} variables to {output_path}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 2 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython2", + "version": "2.7.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index 4e39772..be5fd24 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "peft>=0.18.1", "llm2vec", "setuptools<81", - "geotessera>=0.8.0", + "geotessera>=0.9.0", "ipykernel>=7.2.0", "xgboost>=3.2.0", "albumentations", @@ -44,12 +44,13 @@ dependencies = [ create-data = [ "geemap>=0.36.6", "pipreqs>=0.5.0", - "earthaccess>=0.11.0", +# "earthaccess>=0.11.0", #updated geotessera is incompatible with this one "pyhdf>=0.11.4", ] visualisation = [ "folium>=0.20.0", + "seaborn>=0.13.2", ] [tool.pytest.ini_options] diff --git a/src/data/base_caption_builder.py b/src/data/base_caption_builder.py index 79de05b..feffef2 100644 --- a/src/data/base_caption_builder.py +++ b/src/data/base_caption_builder.py @@ -3,23 +3,33 @@ import random import re from abc import ABC, abstractmethod -from typing import Any, Dict, List, final +from typing import Any, Dict, List, Tuple, final import torch from src.data.base_dataset import BaseDataset +from src.utils.errors import IllegalArgumentCombination class BaseCaptionBuilder(ABC): def __init__( - self, templates_fname: str, concepts_fname: str, data_dir: str, seed: int + self, + templates_fname: str, + concepts_fname: str, + data_dir: str, + seed: int, + n_captions_for_validation: int | str = "all", + return_aux_ids: bool = False, ) -> None: """Interface of caption builder class for converting numerical auxiliary data values into textual descriptions from provided caption templates. :param templates_fname: path to a json file with caption templates. + :param concepts_fname: path to a json file with concepts. :param data_dir: directory where data is stored. :param seed: random seed. + :param n_captions_for_validation: number of captions to randomly sample for validation + :param return_aux_ids: whether to return auxiliary column ids. """ self.data_dir = data_dir @@ -39,6 +49,17 @@ def __init__( self.seed = seed random.seed(self.seed) + if n_captions_for_validation == "all": + self.n = self.__len__ + elif n_captions_for_validation > len(self): + raise IllegalArgumentCombination( + f"Requested {n_captions_for_validation} captions exceeds template dictionary size" + ) + else: + self.n = n_captions_for_validation + + self.return_aux_ids = return_aux_ids + @final def __len__(self): """Number of caption templates.""" @@ -89,43 +110,86 @@ def _build_from_template( """Build caption text from template and row of auxiliary data.""" pass - def random(self, aux_values) -> List[str]: - """Return a caption from a randomly sampled template for each data point.""" - formatted_rows = [] + def random(self, aux_values) -> Tuple[List[str], List[int] | None]: + """Return a caption per location from a randomly sampled template for each data point. + :param aux_values: a batch of auxiliary values to use for random sampling. + :return: a batch of text captions and optionally aux col ids used for each of the caption. + """ batch_size = len(aux_values["aux"]) - template_ids = random.choices( - range(len(self.templates)), - k=batch_size, - ) - for ( - i, - template_idx, - ) in enumerate(template_ids): - row_aux = aux_values["aux"][i] - row_top = aux_values.get("top")[i] if aux_values.get("top") else None - formatted_rows.append( - self._build_from_template(template_idx, aux=row_aux, top=row_top) - ) + # Location captions holders + formatted_location_captions = [] - return formatted_rows + # Ids of used aux col ids per template (location) + if self.return_aux_ids: + ids = [] - def all(self, aux_values) -> List[str]: - """Return a list of captions from all available templates.""" - formatted_rows = [] - for i in range(0, len(aux_values["aux"])): - descriptions = [] + # Sample templates + template_ids = random.choices(range(len(self.templates)), k=batch_size) + for i, template_idx in enumerate(template_ids): + # Get aux and top values per location row_aux = aux_values["aux"][i] row_top = aux_values.get("top")[i] if aux_values.get("top") else None - for template_idx in range(0, len(self)): - descriptions.append( - self._build_from_template(template_idx, aux=row_aux, top=row_top) + # Get filled in template for location + if self.return_aux_ids: + filled_template, template_ids = self._build_from_template( + template_idx, aux=row_aux, top=row_top ) - formatted_rows.append(descriptions) + ids.append(template_ids) + else: + filled_template = self._build_from_template(template_idx, aux=row_aux, top=row_top) + formatted_location_captions.append(filled_template) - return formatted_rows + if self.return_aux_ids: + return formatted_location_captions, ids + return formatted_location_captions + + def sample_multiple_or_all(self, aux_values) -> Tuple[List[str], List[int] | None]: + """Return self.n captions from randomly sampled templates for each data point. + + :param aux_values: a batch of auxiliary values to use for random sampling. + :return: a batch of text captions and optionally aux col ids used for each of the caption. + """ + batch_size = len(aux_values["aux"]) + + # Location captions holders + formatted_location_captions = [] + + # Ids of used aux col ids per template (location) + if self.return_aux_ids: + ids = [] + + for i in range(0, batch_size): + # Get aux and top values per location + row_aux = aux_values["aux"][i] + row_top = aux_values.get("top")[i] if aux_values.get("top") else None + + # Sample templates + template_ids = random.choices(range(len(self.templates)), k=self.n) + + # Get filled in templates for location + filled_in_location_templates = [] + ids_per_location = [] + for template_idx in template_ids: + if self.return_aux_ids: + filled_template, template_ids = self._build_from_template( + template_idx, aux=row_aux, top=row_top + ) + ids_per_location.extend(filled_template) + else: + filled_template = self._build_from_template( + template_idx, aux=row_aux, top=row_top + ) + filled_in_location_templates.append(filled_template) + + if self.return_aux_ids: + ids.append(template_ids) + formatted_location_captions.append(filled_in_location_templates) + if self.return_aux_ids: + return formatted_location_captions, ids + return formatted_location_captions def sync_concepts(self) -> List[str]: for concept in self.concepts: @@ -169,24 +233,25 @@ def get_adjective_for_percentage(value: float) -> str: else: return "almost entirely" + def sample_adjective_for_percentage(percent: float) -> str: """Convert a percentage (0-100) to a descriptive adjective, randomly sampled from synonyms.""" if not 0 <= percent <= 100: raise ValueError(f"Percentage must be between 0 and 100, got {percent}") synonyms = { - "none": ["none", "zero", "absent", "nonexistent"], - "negligible": ["negligible", "trivial", "trace", "barely any", "scarcely any"], - "minimal": ["minimal", "tiny", "very little", "marginal", "meager"], - "slight": ["slight", "small", "modest", "limited", "faint"], - "some": ["some", "a bit of", "a portion of", "partial", "a measure of"], - "moderate": ["moderate", "fair", "reasonable", "middling", "decent"], + "none": ["none", "zero", "absent", "nonexistent"], + "negligible": ["negligible", "trivial", "trace", "barely any", "scarcely any"], + "minimal": ["minimal", "tiny", "very little", "marginal", "meager"], + "slight": ["slight", "small", "modest", "limited", "faint"], + "some": ["some", "a bit of", "a portion of", "partial", "a measure of"], + "moderate": ["moderate", "fair", "reasonable", "middling", "decent"], "considerable": ["considerable", "notable", "meaningful", "appreciable", "marked"], - "substantial": ["substantial", "solid", "sizable", "hefty", "goodly"], - "significant": ["significant", "large", "strong", "pronounced", "prominent"], - "major": ["major", "great", "high", "intense", "serious"], - "extensive": ["extensive", "vast", "sweeping", "far-reaching", "immense"], - "complete": ["complete", "total", "full", "entire", "absolute"], + "substantial": ["substantial", "solid", "sizable", "hefty", "goodly"], + "significant": ["significant", "large", "strong", "pronounced", "prominent"], + "major": ["major", "great", "high", "intense", "serious"], + "extensive": ["extensive", "vast", "sweeping", "far-reaching", "immense"], + "complete": ["complete", "total", "full", "entire", "absolute"], } if percent == 0: @@ -214,4 +279,4 @@ def sample_adjective_for_percentage(percent: float) -> str: else: key = "complete" - return random.choice(synonyms[key]) \ No newline at end of file + return random.choice(synonyms[key]) diff --git a/src/data/base_datamodule.py b/src/data/base_datamodule.py index 7f969cf..04d6a69 100644 --- a/src/data/base_datamodule.py +++ b/src/data/base_datamodule.py @@ -2,7 +2,7 @@ import os import time from functools import partial -from typing import Any, Tuple +from typing import Any, Dict, List, Tuple import numpy as np import pandas as pd @@ -66,6 +66,7 @@ def __init__( self.caption_builder = caption_builder self.caption_builder.sync_with_dataset(self.dataset) self.concept_configs = caption_builder.concepts + self._setup_flag = False @property def tabular_dim(self): @@ -82,9 +83,11 @@ def setup(self, stage: str = "fit") -> None: Called by model trainer (trainer.fit()). """ - # Set up the dataset (download requested modalities) - self.dataset.setup() - self.split_data() + if not self._setup_flag: + # Set up the dataset (download requested modalities) + self.dataset.setup() + self.split_data() + self._setup_flag = True @property def batch_size_per_device(self) -> None: @@ -272,10 +275,13 @@ def split_data(self) -> None: if split_data_from_inds: self.data_train = torch.utils.data.Subset(self.dataset, train_indices) + print(f"Train dataset split size: {len(self.data_train)}") self.data_val = torch.utils.data.Subset(self.dataset, val_indices) + print(f"Validate dataset split size: {len(self.data_val)}") if test_indices is not None: self.data_test = torch.utils.data.Subset(self.dataset, test_indices) + print(f"Test dataset split size: {len(self.data_test)}") else: self.data_test = None diff --git a/src/data/base_dataset.py b/src/data/base_dataset.py index 64ccd7b..1a5d37e 100644 --- a/src/data/base_dataset.py +++ b/src/data/base_dataset.py @@ -9,6 +9,7 @@ import torch from torch.utils.data import Dataset +from src.data_preprocessing.tessera_embeds import NoTileError, PartialTileError from src.utils.data_utils import center_crop_npy TORCH_DTYPES = { @@ -36,6 +37,7 @@ def __init__( use_features: bool = True, csv_name: str = None, dtype: str = "float32", + return_name_loc: bool = False, ) -> None: """Interface for any use case dataset. @@ -143,6 +145,7 @@ def __init__( self.dataset_name: str = dataset_name + "_" + "_".join(modalities) self.mode: str = mode # 'train', 'val', 'test' self.records: dict[str, Any] = self.get_records() + self.return_name_loc: bool = return_name_loc @final def get_records(self) -> dict[str, Any]: @@ -156,6 +159,16 @@ def get_records(self) -> dict[str, Any]: for modality, params in self.modalities.items(): if modality == "coords": columns.extend(["lat", "lon"]) + elif modality in ["aef_avr", "tessera_avr"]: + df = pd.read_csv(params.get("path", KeyError)) + df["name_loc"] = df["name_loc"].astype(str) + + self.df["name_loc"] = self.df["name_loc"].astype(str) + + self.df = self.df.merge(df, on="name_loc", how="left") + + max_no = 128 if modality == "tessera_avr" else 64 + columns.extend([f"emb_{i}" for i in range(0, max_no)]) else: # Add paths self.add_modality_paths_to_df( @@ -258,6 +271,7 @@ def setup_tessera(self) -> None: size = self.modalities["tessera"].get( "size", KeyError('Missing parameter "size" for Tessera modality') ) + version = self.modalities["tessera"].get("version") or "v1.0" # If data does not exist or is empty → full download if not os.path.exists(dst_dir) or len(os.listdir(dst_dir)) == 0: @@ -270,6 +284,7 @@ def setup_tessera(self) -> None: year=year, tile_size=size, cache_dir=self.cache_dir, + version=version, ) # TODO: if we compile the dataset and use zenodo (or sth else) then change to pooch downloading/loading @@ -306,7 +321,7 @@ def setup_tessera(self) -> None: tessera_con=gt, ) continue - except Exception as e: + except NoTileError or PartialTileError as e: print(f"Tile for {fname} could not be retrieved. Error: {e}") else: self.records.pop(i) diff --git a/src/data/butterfly_caption_builder.py b/src/data/butterfly_caption_builder.py index af23298..a4cc22e 100644 --- a/src/data/butterfly_caption_builder.py +++ b/src/data/butterfly_caption_builder.py @@ -1,12 +1,11 @@ import os -from typing import Any, List, override +from typing import List, override import pandas as pd import torch from src.data.base_caption_builder import ( BaseCaptionBuilder, - get_adjective_for_percentage, sample_adjective_for_percentage, ) from src.data.base_dataset import BaseDataset @@ -18,9 +17,22 @@ class ButterflyCaptionBuilder(BaseCaptionBuilder): def __init__( - self, templates_fname: str, concepts_fname: str, data_dir: str, seed: int + self, + templates_fname: str, + concepts_fname: str, + data_dir: str, + seed: int, + n_captions_for_validation: int | str = "all", + return_aux_ids: bool = False, ) -> None: - super().__init__(templates_fname, concepts_fname, data_dir, seed) + super().__init__( + templates_fname, + concepts_fname, + data_dir, + seed, + n_captions_for_validation, + return_aux_ids, + ) @override def sync_with_dataset(self, dataset: BaseDataset) -> None: @@ -44,7 +56,6 @@ def sync_with_dataset(self, dataset: BaseDataset) -> None: "description": description, "units": units, } - self.sync_concepts() def get_corine_column_keys(self): @@ -113,6 +124,9 @@ def _build_from_template( template = self.templates[template_idx] tokens = self.tokens_in_template[template_idx] replacements = {} + if self.return_aux_ids: + ids = [] + for token in tokens: init_token = token if "top" in token: @@ -126,6 +140,8 @@ def _build_from_template( ) idx = values_dict["id"] + if self.return_aux_ids: + ids.append(idx) value = aux[idx].item() formatted_desc = values_dict["description"].lower() or "" @@ -143,6 +159,8 @@ def _build_from_template( replacements[init_token] = formatted_desc template = self._fill(template, replacements) + if self.return_aux_ids: + return template, ids return template diff --git a/src/data/butterfly_dataset.py b/src/data/butterfly_dataset.py index aa6f342..ce49704 100644 --- a/src/data/butterfly_dataset.py +++ b/src/data/butterfly_dataset.py @@ -1,5 +1,5 @@ import os -from typing import Any, Dict, List, override +from typing import Any, Dict, override import numpy as np import pooch @@ -23,6 +23,7 @@ def __init__( cache_dir: str = None, mock: bool = False, dtype: str = "float32", + return_name_loc: bool = False, ) -> None: """A dataset implementation for the Butterfly diversity use case. @@ -40,7 +41,7 @@ def __init__( assert not ( use_unlabelled_data and use_target_data - ), "Joint use of unlabelled and target data is not supported." + ), "Joint use of unlabelled and target data is not supported yet." if use_unlabelled_data: dataset_name = ["s2bms", "s2bms-unlabelled-20260529"] else: @@ -54,9 +55,10 @@ def __init__( dataset_name=dataset_name, seed=seed, cache_dir=cache_dir, - implemented_mod={"s2", "tessera", "coords", "aef"}, + implemented_mod={"s2", "tessera", "coords", "aef", "aef_avr", "tessera_avr"}, mock=mock, dtype=dtype, + return_name_loc=return_name_loc, ) def setup(self): @@ -192,6 +194,16 @@ def __getitem__(self, idx: int) -> Dict[str, Any]: formatted_row["eo"][modality] = self.load_tessera(row["tessera_path"]) elif modality == "aef": formatted_row["eo"][modality] = self.load_aef(row["aef_path"]) + elif modality == "aef_avr": + formatted_row["eo"][modality] = torch.tensor( + [row[f"emb_{i}"] for i in range(64)], + dtype=getattr(torch, self.modalities["aef_avr"].get("dtype")), + ) + elif modality == "tessera_avr": + formatted_row["eo"][modality] = torch.tensor( + [row[f"emb_{i}"] for i in range(128)], + dtype=getattr(torch, self.modalities["tessera_avr"].get("dtype")), + ) if self.use_target_data: formatted_row["target"] = torch.tensor( @@ -208,6 +220,9 @@ def __getitem__(self, idx: int) -> Dict[str, Any]: else: formatted_row["aux"][aux_cat] = [row[v] for v in vals] + if self.return_name_loc: + formatted_row["name_loc"] = row["name_loc"] + return formatted_row diff --git a/src/data/collate_fns.py b/src/data/collate_fns.py index 2c001b4..b5f3739 100644 --- a/src/data/collate_fns.py +++ b/src/data/collate_fns.py @@ -36,10 +36,18 @@ def collate_fn( if batch[0].get("target") is not None: batch_collected["target"] = smart_stack([item["target"] for item in batch]) + if batch[0].get("name_loc") is not None: + batch_collected["name_loc"] = smart_stack([item["name_loc"] for item in batch]) + # convert aux into captions if mode == "train": batch_collected["text"] = caption_builder.random(batch_collected["aux"]) else: - batch_collected["text"] = caption_builder.all(batch_collected["aux"]) + batch_collected["text"] = caption_builder.sample_multiple_or_all(batch_collected["aux"]) + + # If requested to return aux_ids, recompile stacks + if caption_builder.return_aux_ids: + batch_collected["text_aux_ids"] = batch_collected["text"][1] + batch_collected["text"] = batch_collected["text"][0] return batch_collected diff --git a/src/data_preprocessing/tessera_embeds.py b/src/data_preprocessing/tessera_embeds.py index 1f9432a..427be38 100644 --- a/src/data_preprocessing/tessera_embeds.py +++ b/src/data_preprocessing/tessera_embeds.py @@ -130,6 +130,7 @@ def get_tessera_embeds( save_dir: str, tile_size: int, tessera_con: GeoTessera | None, + version: str, padding: int = 100, save_tif=False, ) -> None: @@ -236,7 +237,15 @@ def _close_all(): if padding > 500: raise NoTileError(f"Padding {padding} > 500") return get_tessera_embeds( - lon, lat, name_loc, year, save_dir, tile_size, tessera_con, padding=padding + 100 + lon, + lat, + name_loc, + year, + save_dir, + tile_size, + tessera_con, + padding=padding + 100, + version=version, ) crop = mosaic[row_min:row_max, col_min:col_max, :] @@ -276,7 +285,14 @@ def _close_all(): # Log its metadata meta_df = pd.DataFrame( - {"id": [name_loc], "year": [year], "lon": [lon], "lat": [lat], "crs": [utm_crs]} + { + "id": [name_loc], + "year": [year], + "lon": [lon], + "lat": [lat], + "crs": [utm_crs], + "version": [version], + }, ) meta_file = f"{save_dir}/meta.csv" @@ -298,6 +314,7 @@ def tessera_from_df( tile_size: int = 256, cache_dir: str = "temp/", logs_dir: str = "logs", + version="v1", ) -> None: """Obtains Tessera embeddings from a CSV file for each (lon, lat). @@ -306,19 +323,29 @@ def tessera_from_df( :param year: year for the embeddings :param tile_size: tile size in meters :param cache_dir: path to cache directory + :param version: version of the tessera dataset :return: None """ - + assert version in ["v1", "v1.1"] + dataset_var = {"v1.1": "cambridge", "1.1": "cambridge", "1.0": "vultr", "v1.0": "vultr"}[ + version + ] # Tessera connection cache_dir = os.path.join(cache_dir, "tessera") - gt = GeoTessera(cache_dir=cache_dir, embeddings_dir=cache_dir, dataset_version="v1") - + gt = GeoTessera( + cache_dir=cache_dir, + embeddings_dir=cache_dir, + dataset_version=version, + dataset_variant=dataset_var, + ) # Iter each coord n = len(model_ready_df) for i, row in model_ready_df.iterrows(): print(f"{i}/{n}") try: - get_tessera_embeds(row.lon, row.lat, row.name_loc, year, f"{data_dir}/", tile_size, gt) + get_tessera_embeds( + row.lon, row.lat, row.name_loc, year, f"{data_dir}/", tile_size, gt, version + ) except Exception as e: if isinstance(e, NoTileError): path = os.path.join(logs_dir, "tessera_skipped.txt") @@ -389,24 +416,29 @@ def inspect_np_arr_as_tiff( if __name__ == "__main__": - os.chdir("../..") + # os.chdir("../..") print(os.getcwd()) # df = pd.read_csv("data/heat_guatemala/model_ready_heat_guatemala.csv") # df = pd.read_csv("/lustre/backup/SHARED/AIN/aether/data/s2bms/model_ready_s2bms.csv") df = pd.read_csv("data/s2bms/model_ready_s2bms.csv") - # df.sort_values(by="name_loc", inplace=True, ascending=False) + + # df = pd.read_csv("data/s2bms/model_ready_s2bms-unlabelled-20260529.csv") + # df = df[df.split == "train"] + df.sort_values(by="name_loc", inplace=True, ascending=False) + df.reset_index(drop=True, inplace=True) + if os.path.exists("logs/tessera_skipped.txt"): with open(os.path.join("logs", "tessera_skipped.txt")) as f: skipped = set(f.read().splitlines()) df = df[~df.name_loc.isin(skipped)] - # df.sort_values('name_loc', ascending=False, inplace=True) tessera_from_df( df, "data/s2bms/eo/tessera", - year=2024, + year=2019, tile_size=256, cache_dir="data/cache", + version="v1.1", ) diff --git a/src/data_preprocessing/yield_africa_tessera_preprocess.py b/src/data_preprocessing/yield_africa_tessera_preprocess.py index 3a1e770..cea8a85 100644 --- a/src/data_preprocessing/yield_africa_tessera_preprocess.py +++ b/src/data_preprocessing/yield_africa_tessera_preprocess.py @@ -36,8 +36,8 @@ import multiprocessing import os import socket -import time import sys +import time from pathlib import Path # Ensure the project root is on sys.path when the script is run directly. @@ -45,7 +45,11 @@ import pandas as pd -from src.data_preprocessing.tessera_embeds import NoTileError, PartialTileError, get_tessera_embeds +from src.data_preprocessing.tessera_embeds import ( + NoTileError, + PartialTileError, + get_tessera_embeds, +) log = logging.getLogger(__name__) @@ -65,6 +69,7 @@ def _init_worker(cache_dir: str, use_local_registry: bool, registry_dir: str) -> None: """Pool initializer: runs once per worker process to set up GeoTessera.""" from geotessera import GeoTessera + global _process_gt socket.setdefaulttimeout(60) embeddings_dir = str(Path(cache_dir) / "raw") @@ -78,13 +83,18 @@ def _init_worker(cache_dir: str, use_local_registry: bool, registry_dir: str) -> def _worker_fetch(args: tuple) -> str: """Multiprocessing worker — reuses the per-process GeoTessera instance. - Must be a top-level function so it is picklable across processes. - Returning normally means success; raising means error (logged by caller). + Must be a top-level function so it is picklable across processes. Returning normally means + success; raising means error (logged by caller). """ lon, lat, name_loc, year, save_dir, tile_size = args get_tessera_embeds( - lon=lon, lat=lat, name_loc=name_loc, year=year, - save_dir=save_dir, tile_size=tile_size, tessera_con=_process_gt, + lon=lon, + lat=lat, + name_loc=name_loc, + year=year, + save_dir=save_dir, + tile_size=tile_size, + tessera_con=_process_gt, ) return name_loc @@ -160,8 +170,8 @@ def fetch_tessera_tiles( # Per-task timeout: when a worker process exceeds this, it is killed and # the record added to stuck.txt. Multiprocessing (unlike threading) allows # true process termination, so stuck downloads cannot block forward progress. - HEARTBEAT = 15 # seconds between "still fetching" log lines - TILE_TIMEOUT = 60 # seconds per record before the worker process is killed + HEARTBEAT = 15 # seconds between "still fetching" log lines + TILE_TIMEOUT = 60 # seconds per record before the worker process is killed # Note: stuck workers spin at 100% CPU and leak ~55 MB/s of rasterio # MemoryFile objects inside GeoTessera's fetch loop. 60s caps peak memory # at ~3 GB per stuck record and recovers 3x faster than the old 180s limit. @@ -177,13 +187,17 @@ def fetch_tessera_tiles( else: stuck_records = set(stuck_file.read_text().splitlines()) if stuck_records: - print(f" Skipping {len(stuck_records)} previously-stuck record(s): {sorted(stuck_records)}") + print( + f" Skipping {len(stuck_records)} previously-stuck record(s): {sorted(stuck_records)}" + ) rows = [row for _, row in df.iterrows() if row.name_loc not in stuck_records] done = 0 _pool_initargs = (cache_dir, _use_local_registry, str(_default_registry_dir)) - pool = multiprocessing.Pool(processes=workers, initializer=_init_worker, initargs=_pool_initargs) + pool = multiprocessing.Pool( + processes=workers, initializer=_init_worker, initargs=_pool_initargs + ) try: for row in rows: args = (row.lon, row.lat, row.name_loc, int(row.year), str(save_dir), tile_size) @@ -204,7 +218,9 @@ def fetch_tessera_tiles( print(f" Skipped {row.name_loc}: no TESSERA data for this location/year") break except PartialTileError: - print(f" Skipped {row.name_loc}: tile too close to mosaic edge, not enough context") + print( + f" Skipped {row.name_loc}: tile too close to mosaic edge, not enough context" + ) break except Exception as exc: print(f" ERROR fetching {row.name_loc}: {exc}") @@ -213,7 +229,9 @@ def fetch_tessera_tiles( if timed_out: pool.terminate() pool.join() - pool = multiprocessing.Pool(processes=workers, initializer=_init_worker, initargs=_pool_initargs) + pool = multiprocessing.Pool( + processes=workers, initializer=_init_worker, initargs=_pool_initargs + ) with open(stuck_file, "a") as fh: fh.write(row.name_loc + "\n") print( diff --git a/src/inference.py b/src/inference.py index dc15b50..6de9d6c 100644 --- a/src/inference.py +++ b/src/inference.py @@ -32,12 +32,27 @@ def main(cfg: DictConfig) -> Optional[float]: # If a merged inference ckpt is provided, just load it. inference_ckpt_path = cfg.get("inference_ckpt_path") if inference_ckpt_path: - model = load_inference_model(inference_ckpt_path) + model = load_inference_model(inference_ckpt_path, cfg.paths.cache_dir) # Otherwise merge model from two checkpoints else: model = merge_inference_model(cfg, save_ckpt=True) + model.to("mps") # TODO: do what you need with the inference model + if cfg.data: + datamodule = hydra.utils.instantiate(cfg.get("data")) + datamodule.setup() + + concepts = [ + c["concept_caption"] for c in datamodule.caption_builder.__dict__["concepts"] + ] # or other source of concepts + text_embeds = model.forward_text(concepts) + + # per location (batching uses location text generation) + for d in datamodule.data_train: + b = {"eo": {"aef_avr": d["eo"]["aef_avr"].unsqueeze(0).to("mps")}} + geo_embeds, pred = model.forward_geo(b) + # TODO: do what you need with the inference model return diff --git a/src/models/base_model.py b/src/models/base_model.py index d0f92e8..19694cb 100644 --- a/src/models/base_model.py +++ b/src/models/base_model.py @@ -45,7 +45,9 @@ def __init__( self.save_hyperparameters( ignore=[ "geo_encoder", + "geo_adapter", "text_encoder", + "text_adapter", "prediction_head", "optimizer", "scheduler", @@ -78,7 +80,7 @@ def setup(self, stage: str) -> None: """Updates model based data-bound configurations (through datamodule), This method is called after trainer is initialized and datamodule is available.""" if self.setup_flag: - print(f"Model {self.__str__()} is already set up!") + print("Model is already set up!") return # If trainer is attached get num_classes and tabular_dim from datamodule (data-dependent) @@ -86,6 +88,9 @@ def setup(self, stage: str) -> None: self.num_classes = self.trainer.datamodule.num_classes self.tabular_dim = self.trainer.datamodule.tabular_dim + # set up loss if needed + self.loss_fn.setup(datamodule=self.trainer.datamodule, device=self.device) + # Per model logic of setting up self._setup(stage) self.setup_flag = True @@ -104,8 +109,11 @@ def _setup(self, stage: str) -> None: @final def full_freezer(self): """Freeze the whole network.""" + print("--------Frozen--------") for name, param in self.named_parameters(): param.requires_grad = False + print("Full model") + print("------------------------") for name, module in self.named_modules(): module.eval() @@ -215,12 +223,20 @@ def update_configs(self, cfg): """Update hyper-parameters from the model.""" if hasattr(self, "geo_encoder"): self.geo_encoder.update_configs(cfg["geo_encoder"]) + if hasattr(self, "geo_adapter") and self.geo_adapter is not None: + self.geo_adapter.cfg_dict = cfg["geo_adapter"] if hasattr(self, "text_encoder"): self.text_encoder.cfg_dict = cfg["text_encoder"] + if hasattr(self, "text_adapter") and self.text_adapter is not None: + self.text_adapter.cfg_dict = cfg["text_adapter"] - if hasattr(self, "prediction_head"): - self.prediction_head.cfg_dict = cfg["prediction_head"] + if ( + hasattr(self, "prediction_head") + and self.prediction_head + and cfg.get("prediction_head") + ): + self.prediction_head.update_configs(cfg["prediction_head"]) def on_save_checkpoint(self, checkpoint): """Save checkpoint. @@ -258,13 +274,18 @@ def on_save_checkpoint(self, checkpoint): } ) - if hasattr(self, "geo_encoder"): + if hasattr(self, "geo_encoder") and self.geo_encoder is not None: checkpoint["hyper_parameters"]["geo_encoder"] = self.geo_encoder.cfg_dict - if hasattr(self, "prediction_head"): + if hasattr(self, "geo_adapter") and self.geo_adapter is not None: + checkpoint["hyper_parameters"]["geo_adapter"] = self.geo_adapter.cfg_dict + + if hasattr(self, "prediction_head") and self.prediction_head is not None: checkpoint["hyper_parameters"]["prediction_head"] = self.prediction_head.cfg_dict - if hasattr(self, "text_encoder"): - checkpoint["hyper_parameters"]["text_encoder"] = self.text_encoder.cfg_dict + if hasattr(self, "text_encoder") and self.text_encoder is not None: + checkpoint["hyper_parameters"]["text_encoder"] = self.text_encoder.cfg_dict + if hasattr(self, "text_adapter") and self.text_adapter is not None: + checkpoint["hyper_parameters"]["text_adapter"] = self.text_adapter.cfg_dict return def on_load_checkpoint(self, checkpoint): diff --git a/src/models/components/geo_encoders/adopt_encoder.py b/src/models/components/geo_encoders/adopt_encoder.py index 2a5bd88..969be07 100644 --- a/src/models/components/geo_encoders/adopt_encoder.py +++ b/src/models/components/geo_encoders/adopt_encoder.py @@ -1,5 +1,3 @@ -from typing import Dict, List, override - import hydra import torch @@ -19,18 +17,28 @@ def adopt_encoder(ckpt_path: str) -> BaseGeoEncoder: # Get skeleton geo_config = ckpt["hyper_parameters"].get("geo_encoder") encoder: BaseGeoEncoder = hydra.utils.instantiate(geo_config) + print("---Adopted encoder------") + encoder.adopted = True encoder.setup() encoder.cfg_dict = geo_config - print("------------------------") + encoder.cfg_dict.update( + { + "adopted": True, + "save_weights": bool( + "geo_encoder" in "".join(ckpt["hyper_parameters"].get("trainable_modules", [])) + ), + } + ) # Load in the weights state_dict = { - k.replace("geo_encoder.", ""): v + k.replace("geo_encoder.", "", 1): v for k, v in ckpt["state_dict"].items() if "geo_encoder." in k } res = encoder.load_state_dict(state_dict, strict=False) - log_model_loading("geo_encoder_ckpt", res) + log_model_loading("geo_encoder_ckpt", res, True) + print("------------------------") return encoder diff --git a/src/models/components/geo_encoders/average_encoder.py b/src/models/components/geo_encoders/average_encoder.py index f3152bd..46258d5 100644 --- a/src/models/components/geo_encoders/average_encoder.py +++ b/src/models/components/geo_encoders/average_encoder.py @@ -36,7 +36,7 @@ def _setup(self) -> List[str]: def forward(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: """Data forward pass through the encoder.""" tile = batch.get("eo", {}).get(self.geo_data_name) - + # NaNs vs 0s are specified in dataset configs and returned from BaseDataset class feats = self.geo_encoder(tile.nanmean(dim=(-2, -1))) if self.extra_projector: diff --git a/src/models/components/geo_encoders/base_geo_encoder.py b/src/models/components/geo_encoders/base_geo_encoder.py index 8711140..2b6062b 100644 --- a/src/models/components/geo_encoders/base_geo_encoder.py +++ b/src/models/components/geo_encoders/base_geo_encoder.py @@ -19,7 +19,9 @@ def __init__(self) -> None: self.allowed_geo_data_names: list[str] | None = None self.geo_data_name: str | None = None + self.adopted: bool = False + @final def update_configs(self, cfg): if len(self.cfg_dict) == 0: self.cfg_dict = cfg @@ -41,6 +43,8 @@ def setup(self, verbose=1) -> List[str]: if verbose > 0: print(f"Model set up with {self.__str__()}") self.setup_flag = True + if self.adopted: + trainable_modules = [] return trainable_modules @abstractmethod diff --git a/src/models/components/geo_encoders/cnn_encoder.py b/src/models/components/geo_encoders/cnn_encoder.py index 918b755..75c80c5 100644 --- a/src/models/components/geo_encoders/cnn_encoder.py +++ b/src/models/components/geo_encoders/cnn_encoder.py @@ -3,16 +3,13 @@ import torch import torchvision.models as models from torch import nn -from torchgeo.models import resnet50, ResNet50_Weights, ResNet18_Weights, resnet18 +from torchgeo.models import ResNet18_Weights, ResNet50_Weights, resnet18, resnet50 from src.models.components.geo_encoders.base_geo_encoder import BaseGeoEncoder from src.utils.errors import IllegalArgumentCombination -RN_DIM = { - 18 : 512, - 34: 512, - 50: 2048 -} +RN_DIM = {18: 512, 34: 512, 50: 2048} + class CNNEncoder(BaseGeoEncoder): """Convolutional neural network EO encoder. Adapted from PECL. @@ -41,7 +38,12 @@ def __init__( assert resnet_version in [18, 34, 50], f"Unsupported resnet version: {resnet_version}" self.resnet_version = resnet_version - assert pretrained_cnn in ["imagenet", "IMAGENET1K_V1", 'SSL4EO_RGB_MOCO', None], f"Unsupported pretrained_cnn: {pretrained_cnn}" + assert pretrained_cnn in [ + "imagenet", + "IMAGENET1K_V1", + "SSL4EO_RGB_MOCO", + None, + ], f"Unsupported pretrained_cnn: {pretrained_cnn}" self.pretrained_cnn = pretrained_cnn self.output_dim = RN_DIM[resnet_version] @@ -52,7 +54,9 @@ def __init__( self.geo_data_name = geo_data_name self.set_n_input_bands(input_n_bands) - assert (self.input_n_bands >= 3 and type(self.input_n_bands) is int), f"input_n_bands must be int >=3, got {self.input_n_bands}" + assert ( + self.input_n_bands >= 3 and type(self.input_n_bands) is int + ), f"input_n_bands must be int >=3, got {self.input_n_bands}" def set_n_input_bands(self, n_bands: int | None = None) -> None: """Sets number of input bands based on geo_data_name if n_bands is None. @@ -77,6 +81,7 @@ def set_n_input_bands(self, n_bands: int | None = None) -> None: @override def _setup(self) -> List[str]: """Gets backbone model given configuration stored in self. + :return: backbone model """ trainable_modules = [] @@ -87,7 +92,9 @@ def _setup(self) -> List[str]: if self.resnet_version == 18: self.geo_encoder = resnet18(weights=ResNet18_Weights.SENTINEL2_RGB_MOCO) elif self.resnet_version == 34: - raise IllegalArgumentCombination('SSL4EO_RGB_MOCO weights are not available for RN-34') + raise IllegalArgumentCombination( + "SSL4EO_RGB_MOCO weights are not available for RN-34" + ) else: self.geo_encoder = resnet50(weights=ResNet50_Weights.SENTINEL2_RGB_MOCO) # Imagenet @@ -123,7 +130,7 @@ def _setup(self) -> List[str]: self.geo_encoder.conv1.weight[:, i, :, :] = weight[:, i % 3, :, :] # Ensure replaced layer is not frozen - trainable_modules.append('geo_encoder.conv1') + trainable_modules.append("geo_encoder.conv1") # I think for features fc often is replaced with identity? self.geo_encoder.fc = nn.Identity() @@ -145,8 +152,10 @@ def forward( :param batch: input batch :return: extracted features """ - eo_data = batch.get("eo", KeyError(f"Batch must contain batch['eo']")) - eo_data = eo_data.get(self.geo_data_name, KeyError(f"Batch must contain batch['eo']['{self.geo_data_name}']")) + eo_data = batch.get("eo", KeyError("Batch must contain batch['eo']")) + eo_data = eo_data.get( + self.geo_data_name, KeyError(f"Batch must contain batch['eo']['{self.geo_data_name}']") + ) dtype = self.dtype if eo_data.dtype != dtype: @@ -156,4 +165,4 @@ def forward( if self.extra_projector: feats = self.extra_projector(feats) - return feats \ No newline at end of file + return feats diff --git a/src/models/components/geo_encoders/encoder_wrapper.py b/src/models/components/geo_encoders/encoder_wrapper.py index a44a3fa..d8d467d 100644 --- a/src/models/components/geo_encoders/encoder_wrapper.py +++ b/src/models/components/geo_encoders/encoder_wrapper.py @@ -4,9 +4,7 @@ import torch.nn as nn from src.models.components.geo_encoders.base_geo_encoder import BaseGeoEncoder -from src.models.components.geo_encoders.identity_encoder import IdentityEncoder from src.models.components.geo_encoders.tabular_encoder import TabularEncoder -from src.utils.errors import IllegalArgumentCombination class EncoderWrapper(BaseGeoEncoder): @@ -98,10 +96,6 @@ def _setup(self) -> List[str]: # Configure adapter/projector if requested if "projector" in branch: - if isinstance(encoder, IdentityEncoder): - raise IllegalArgumentCombination( - "Identity encoder cannot have linear projector" - ) projector = branch["projector"] projector.set_input_dim(input_dim=branch_dim) diff --git a/src/models/components/geo_encoders/geoclip.py b/src/models/components/geo_encoders/geoclip.py index 9224d60..85eca81 100644 --- a/src/models/components/geo_encoders/geoclip.py +++ b/src/models/components/geo_encoders/geoclip.py @@ -2,7 +2,6 @@ import torch from geoclip import LocationEncoder -from torch.nn import functional as F from src.models.components.geo_encoders.base_geo_encoder import BaseGeoEncoder diff --git a/src/models/components/geo_encoders/identity_encoder.py b/src/models/components/geo_encoders/identity_encoder.py index 29f9bfe..19fbf4a 100644 --- a/src/models/components/geo_encoders/identity_encoder.py +++ b/src/models/components/geo_encoders/identity_encoder.py @@ -11,24 +11,25 @@ def __init__( self, geo_data_name="aef", ) -> None: - """Encoder to avreage tile values into a 1D vector. + """Encoder to pass through the data with no changes. :param geo_data_name: modality name """ super().__init__() - self.dict_n_bands_default = {"s2": 4, "aef": 64, "tessera": 128} + self.dict_n_bands_default = {"aef": 64, "tessera": 128, "aef_avr": 64, "tessera_avr": 128} self.allowed_geo_data_names: list[str] = list(self.dict_n_bands_default.keys()) assert ( geo_data_name in self.allowed_geo_data_names ), f"geo_data_name must be one of {self.allowed_geo_data_names}, got {geo_data_name}" self.geo_data_name = geo_data_name + self.output_dim = self.dict_n_bands_default[geo_data_name] + @override def _setup(self) -> List[str]: """Configures modules and returns newly initialised, trainable module names.""" - self.output_dim = None self.geo_encoder = nn.Identity() return [] @@ -37,9 +38,6 @@ def forward(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: """Data forward pass through the encoder.""" tile = batch.get("eo", {}).get(self.geo_data_name) - if self.output_dim is None: - self.output_dim = tile.shape - feats = self.geo_encoder(tile) if self.extra_projector: feats = self.extra_projector(feats) @@ -47,4 +45,24 @@ def forward(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: @property def device(self): - return + if self.extra_projector: + devices = {p.device for p in self.extra_projector.parameters()} + if len(devices) > 1: + raise RuntimeError("Extra encoder is on multiple devices") + elif len(devices) == 0: + return None + return devices.pop() + else: + return None + + @property + def dtype(self): + if self.extra_projector: + dtypes = {p.dtype for p in self.extra_projector.parameters()} + if len(dtypes) > 1: + raise RuntimeError("Extra encoder has multiple dtypes") + elif len(dtypes) == 0: + return None + return dtypes.pop() + else: + return None diff --git a/src/models/components/loss_fns/base_loss_fn.py b/src/models/components/loss_fns/base_loss_fn.py index cfa0097..780bb09 100644 --- a/src/models/components/loss_fns/base_loss_fn.py +++ b/src/models/components/loss_fns/base_loss_fn.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Dict +from typing import Any, Dict import torch from torch import nn @@ -21,3 +21,8 @@ def forward( **kwargs, ) -> torch.Tensor: pass + + def setup(self, **kwargs: Any) -> None: + """Setup method for losses in case they need some parameters from the dataset (e.g., for + standardisation in the soft contrastive loss.""" + pass diff --git a/src/models/components/loss_fns/bce_loss.py b/src/models/components/loss_fns/bce_loss.py index b4f6663..d3ddab1 100644 --- a/src/models/components/loss_fns/bce_loss.py +++ b/src/models/components/loss_fns/bce_loss.py @@ -21,9 +21,10 @@ def forward( mode: str | None = None, **kwargs, ) -> torch.Tensor or Dict[str, torch.Tensor]: + """Forward pass to get BCE loss.""" labels = labels if labels is not None else batch.get("target") - loss = self.criterion(pred, labels) + loss = self.criterion(pred, labels.to(pred.dtype)) if "return_label" in kwargs: return {f"{mode}_{self.name}": loss} diff --git a/src/models/components/loss_fns/soft_contrastive_loss.py b/src/models/components/loss_fns/soft_contrastive_loss.py new file mode 100644 index 0000000..758ea83 --- /dev/null +++ b/src/models/components/loss_fns/soft_contrastive_loss.py @@ -0,0 +1,129 @@ +import json +from typing import Dict, List, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing_extensions import override + +from src.data.base_datamodule import BaseDataModule +from src.models.components.loss_fns.base_loss_fn import BaseLossFn + + +class SoftContrastiveLoss(BaseLossFn): + def __init__(self, stats_file: str, temperature: float = 0.07, sigma: float = 1.0): + """Soft-Contrastive Loss Function which allows environmentally similar locations to be + treated as soft positives. + + :param temperature: Soft-Contrastive Loss Function Temperature + :param sigma: Soft-Contrastive Loss Function Strength + :param stats_file: path to a json file with statistics of the aux cols (train split). + """ + super().__init__() + + self.log_temp = nn.Parameter(torch.log(torch.tensor(temperature))) + self.sigma = sigma + + self.stats = json.load(open(stats_file)) + + self.name = "SoftContrastiveLoss" + + @override + def setup(self, datamodule: BaseDataModule, device: torch.device): + """Extract auxiliary value statistics into tensors with correct column id sequence.""" + max_id = len(datamodule.caption_builder.column_to_metadata_map["aux"]) + means = torch.zeros(max_id) + stds = torch.ones(max_id) + + for name, stats in self.stats.items(): + # Synchronise aux col names into proper col ids + idx = datamodule.caption_builder.column_to_metadata_map["aux"][name]["id"] + means[idx] = stats["mean"] + stds[idx] = stats["std"] + + self.means = means + self.means = self.means.to(device) + self.stds = stds + 1e-8 + self.stds = self.stds.to(device) + + @override + def forward( + self, + eo_mod: torch.Tensor, + text_mod: torch.Tensor, + aux_values: torch.Tensor, + aux_ids_per_caption: List[List[int]], + mode: str | None = None, + **kwargs, + ) -> Union[torch.Tensor, Dict[str, torch.Tensor]]: + """Forward computation.""" + # Get target matrix + T = self._get_soft_target_matrix( + aux_values=aux_values, aux_ids_per_caption=aux_ids_per_caption + ) + + # Normalize targets + T_eo2text = T / T.sum(dim=1, keepdim=True) + T_text2eo = T / T.sum(dim=0, keepdim=True) + + # Normalise inputs + eo_mod = F.normalize(eo_mod, dim=-1) + text_mod = F.normalize(text_mod, dim=-1) + + # Clip temperature to not exceed 100 + temperature = torch.clamp(self.log_temp.exp(), max=100) + + # Get cosine similarity + dot_product = (eo_mod @ text_mod.T) / temperature + + # Cross entropy: sum(-target * log_prob) + loss_eo2text = -(T_eo2text * F.log_softmax(dot_product, dim=1)).sum(dim=1).mean() + loss_text2eo = -(T_text2eo * F.log_softmax(dot_product, dim=0)).sum(dim=0).mean() + + loss = (loss_eo2text + loss_text2eo) / 2 + if "return_label" in kwargs: + return {f"{mode}_{self.name}": loss} + else: + return loss + + def _get_soft_target_matrix( + self, + aux_values: torch.Tensor, + aux_ids_per_caption: list[list[int]] | None, + ) -> torch.Tensor: + """Puts together a target matrix based on auxiliary value similarity (either all, or + specific per caption template). + + :param aux_values: auxiliary column values (standardised). + :param aux_ids_per_caption: list of ids of auxiliary columns used per caption template. + :return: target matrix which tells how each row should be similar to column based on aux + values (all or specific ones per caption template). + """ + batch_size, n_aux_cols = aux_values.shape + device = aux_values.device + + # Standardise aux values + if self.stats is not None: + aux_values = (aux_values - self.means) / self.stds + + # Soft targets based on the squared difference between location i and j for all aux values + diffs = (aux_values.unsqueeze(1) - aux_values.unsqueeze(0)) ** 2 + + # Create a mask based on aux_columns used per location caption + if aux_ids_per_caption: + # Create a mask + mask = torch.zeros((batch_size, n_aux_cols), dtype=aux_values.dtype, device=device) + for j, used_cols in enumerate(aux_ids_per_caption): + if used_cols: + mask[j, used_cols] = 1.0 + mask = mask.unsqueeze(0) + + # Calculate the masked average distance for the selected aux cols + dist = (diffs * mask).sum(dim=-1) / (mask.sum(dim=-1) + 1e-8) + else: + # Distance based on all aux columns + dist = diffs.mean(dim=-1) + + # Convert distances to similarities using a Gaussian kernel + T = torch.exp(-dist / (2 * self.sigma**2)) # (N, N) + return T diff --git a/src/models/components/pred_heads/base_pred_head.py b/src/models/components/pred_heads/base_pred_head.py index 02e2af6..60d493b 100644 --- a/src/models/components/pred_heads/base_pred_head.py +++ b/src/models/components/pred_heads/base_pred_head.py @@ -76,4 +76,11 @@ def dtype(self) -> torch.dtype | None: raise RuntimeError("Prediction head has multiple dtypes") elif len(dtypes) == 0: return None - return dtypes.pop() \ No newline at end of file + return dtypes.pop() + + @final + def update_configs(self, cfg): + if len(self.cfg_dict) == 0: + self.cfg_dict = cfg + else: + print("Configs for prediction head not updated") diff --git a/src/models/components/projectors_adapters/__init__.py b/src/models/components/projectors_adapters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/models/components/projectors_adapters/base_encoder.py b/src/models/components/projectors_adapters/base_encoder.py new file mode 100644 index 0000000..cfa7da1 --- /dev/null +++ b/src/models/components/projectors_adapters/base_encoder.py @@ -0,0 +1,65 @@ +from abc import ABC, abstractmethod +from typing import Dict, List, final + +import torch +from torch import nn + + +class BaseEncoder(nn.Module, ABC): + def __init__(self) -> None: + super().__init__() + + # Modules + self.output_dim: int | None = None + self.setup_flag: bool = False + self.cfg_dict: Dict = {} + self.input_dim: int | None = None + self.cfg_dict: Dict = {} + + @final + def set_input_dim(self, input_dim: int) -> None: + self.input_dim = input_dim + + @final + def setup(self, verbose=1) -> List[str]: + """Configures modules. + + Gets called in model.setup() method. Returns names of any new module configured to be added + to the trainable modules list. + """ + if self.setup_flag: + print(f"Module {self.__str__()} is already set up.") + return [] + else: + trainable_modules = self._setup() + if verbose > 0: + print(f"Model set up with {self.__str__()}") + self.setup_flag = True + return trainable_modules + + @abstractmethod + def _setup(self) -> List[str]: + """Configures modules and returns newly initialised, trainable module names.""" + pass + + @abstractmethod + def forward(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: + pass + + @property + def device(self) -> torch.device | None: + devices = {p.device for p in self.parameters()} + if len(devices) > 1: + raise RuntimeError("Encoder is on multiple devices") + elif len(devices) == 0: + return None + return devices.pop() + + @property + def dtype(self) -> torch.dtype | None: + dtypes = {p.dtype for p in self.parameters()} + if len(dtypes) > 1: + raise RuntimeError("Encoder has multiple dtypes") + elif len(dtypes) == 0: + return None + return dtypes.pop() diff --git a/src/models/components/geo_encoders/mlp_projector.py b/src/models/components/projectors_adapters/mlp_projector.py similarity index 81% rename from src/models/components/geo_encoders/mlp_projector.py rename to src/models/components/projectors_adapters/mlp_projector.py index a43316e..1822038 100644 --- a/src/models/components/geo_encoders/mlp_projector.py +++ b/src/models/components/projectors_adapters/mlp_projector.py @@ -3,10 +3,10 @@ import torch from torch import nn -from src.models.components.geo_encoders.base_geo_encoder import BaseGeoEncoder +from src.models.components.projectors_adapters.base_encoder import BaseEncoder -class MLPProjector(BaseGeoEncoder): +class MLPProjector(BaseEncoder): def __init__( self, output_dim: int, @@ -29,9 +29,6 @@ def _setup(self) -> List[str]: self.configure_nn() return ["net"] - def set_input_dim(self, input_dim: int) -> None: - self.input_dim = input_dim - def configure_nn(self) -> None: """Configure the MLP network.""" assert self.input_dim is not None, "input_dim must be defined" @@ -49,6 +46,4 @@ def configure_nn(self) -> None: def forward(self, x: torch.Tensor) -> torch.Tensor: feats = self.net(x) - if self.extra_projector: - feats = self.extra_projector(feats) return feats diff --git a/src/models/components/text_encoders/clip_text_encoder.py b/src/models/components/text_encoders/clip_text_encoder.py index 5052ca5..33254f0 100644 --- a/src/models/components/text_encoders/clip_text_encoder.py +++ b/src/models/components/text_encoders/clip_text_encoder.py @@ -10,7 +10,7 @@ class ClipTextEncoder(BaseTextEncoder): - def __init__(self, hf_cache_dir: str = "../.cache", output_normalization="l2") -> None: + def __init__(self, hf_cache_dir: str = "../.cache", use_geoclip_projector=True) -> None: super().__init__() self.processor = CLIPProcessor.from_pretrained( "openai/clip-vit-large-patch14", @@ -22,15 +22,17 @@ def __init__(self, hf_cache_dir: str = "../.cache", output_normalization="l2") - cache_dir=hf_cache_dir, ) - self.projector = GeoCLIP().image_encoder.mlp + if use_geoclip_projector: + self.projector = GeoCLIP().image_encoder.mlp self.model.vision_model = None self.model.visual_projection = None - self.output_dim = 512 + self.output_dim = 512 if use_geoclip_projector else 768 @override def forward(self, batch: Dict[str, torch.Tensor], mode: str) -> torch.Tensor: + """Forward pass through the text encoder.""" # Get text inputs text_input = batch.get("text") @@ -50,7 +52,11 @@ def forward(self, batch: Dict[str, torch.Tensor], mode: str) -> torch.Tensor: device = next(self.model.parameters()).device text_tokens = {k: v.to(device) for k, v in text_tokens.items()} - text_embeds = self.model.get_text_features(**text_tokens) + if any(p.requires_grad for p in self.model.parameters()): + text_embeds = self.model.get_text_features(**text_tokens) + else: + with torch.no_grad(): + text_embeds = self.model.get_text_features(**text_tokens) # Project if self.projector is not None: diff --git a/src/models/components/text_encoders/llm2clip_text_encoder.py b/src/models/components/text_encoders/llm2clip_text_encoder.py index 0f9af97..eccd667 100644 --- a/src/models/components/text_encoders/llm2clip_text_encoder.py +++ b/src/models/components/text_encoders/llm2clip_text_encoder.py @@ -1,8 +1,8 @@ from typing import Dict, override import torch +from huggingface_hub import snapshot_download from llm2vec import LLM2Vec -from torch.nn import functional as F from transformers import AutoConfig, AutoModel, AutoTokenizer from src.models.components.text_encoders.base_text_encoder import BaseTextEncoder @@ -18,25 +18,26 @@ def __init__(self, hf_cache_dir: str = "../.cache", output_normalization="l2") - :param output_normalization: output normalization type """ super().__init__() + self._download(hf_cache_dir) - # Adapter and image encoder self.projector = AutoModel.from_pretrained( "microsoft/LLM2CLIP-Openai-L-14-224", trust_remote_code=True, dtype=torch.bfloat16, - revision="50ed31c5248d8ff124893719e37829d59376be81", # pin revision for full reproducibility + revision="50ed31c5248d8ff124893719e37829d59376be81", cache_dir=hf_cache_dir, ).eval() - # TODO: If we want to reuse the vision part this is the fix place self.projector.vision_model = None self.projector.visual_projection = None - # The LLM sentence encoder llm_model_name = "microsoft/LLM2CLIP-Llama-3-8B-Instruct-CC-Finetuned" + llm_revision = "main" + config = AutoConfig.from_pretrained( llm_model_name, trust_remote_code=True, + revision=llm_revision, cache_dir=hf_cache_dir, ) config._attn_implementation = "eager" @@ -45,13 +46,17 @@ def __init__(self, hf_cache_dir: str = "../.cache", output_normalization="l2") - llm_model_name, config=config, dtype=torch.bfloat16, - trust_remote_code=False, # local code + trust_remote_code=False, + revision=llm_revision, cache_dir=hf_cache_dir, ) - llm_model.config._name_or_path = ( - "meta-llama/Meta-Llama-3-8B-Instruct" # Workaround for LLM2VEC + llm_model.config._name_or_path = "meta-llama/Meta-Llama-3-8B-Instruct" + + self.processor = AutoTokenizer.from_pretrained( + llm_model_name, + revision=llm_revision, + cache_dir=hf_cache_dir, ) - self.processor = AutoTokenizer.from_pretrained(llm_model_name) # Caption to vector with the llama LLM self.model = LLM2Vec( @@ -60,6 +65,20 @@ def __init__(self, hf_cache_dir: str = "../.cache", output_normalization="l2") - self.output_dim = 1280 + @staticmethod + def _download(hf_cache_dir: str) -> None: + llm_revision = "main" + snapshot_download( + "microsoft/LLM2CLIP-Llama-3-8B-Instruct-CC-Finetuned", + revision=llm_revision, + cache_dir=hf_cache_dir, + ) + snapshot_download( + "microsoft/LLM2CLIP-Openai-L-14-224", + revision="50ed31c5248d8ff124893719e37829d59376be81", + cache_dir=hf_cache_dir, + ) + @override def forward(self, batch: Dict[str, torch.Tensor], mode: str) -> torch.Tensor: """Forward pass through text encoder.""" @@ -67,32 +86,29 @@ def forward(self, batch: Dict[str, torch.Tensor], mode: str) -> torch.Tensor: text_input = batch.get("text") if mode == "train": - text_input = [text_input] - # Embed text and if not train average all templates - avr_embeds = [] - for captions_per_row in text_input: - # LLM is frozen, no gradients needed - with torch.no_grad(): - # Embed - text_embeds = self.model.encode( - captions_per_row, convert_to_tensor=True, device=self.device - ) - - # Change dtype - text_embeds = text_embeds.to( - dtype=self.projector.dtype, device=self.projector.device - ) - - # Project to align with ViT in LLM2CLIP - text_embeds = self.projector.get_text_features(text_embeds) - - if self.extra_projector is not None: - text_embeds = self.extra_projector(text_embeds) - - if mode != "train": - avr_embeds.append(text_embeds.mean(dim=0)) + captions_flattened = text_input + num_per_location = 1 + else: + num_per_location = len(text_input[0]) + captions_flattened = [c for loc in text_input for c in loc] + + # LLM is frozen, no gradients needed + with torch.no_grad(): + # Embed + text_embeds = self.model.encode( + captions_flattened, convert_to_tensor=True, device=self.device + ) + # Change dtype + text_embeds = text_embeds.to(dtype=self.projector.dtype, device=self.projector.device) + + # Project to align with ViT in LLM2CLIP + text_embeds = self.projector.get_text_features(text_embeds) + + if self.extra_projector is not None: + text_embeds = self.extra_projector(text_embeds) if mode != "train": - text_embeds = torch.stack(avr_embeds, dim=0) + text_embeds = text_embeds.reshape(-1, num_per_location, text_embeds.shape[-1]) + text_embeds = text_embeds.mean(dim=1) return text_embeds diff --git a/src/models/inference_model.py b/src/models/inference_model.py index 5c0023e..f07f839 100644 --- a/src/models/inference_model.py +++ b/src/models/inference_model.py @@ -23,6 +23,7 @@ class InferenceModel(BaseModel): def __init__( self, geo_encoder: BaseGeoEncoder | None, + geo_encoder_prediction: BaseGeoEncoder | None, text_encoder: BaseTextEncoder | None, prediction_head: BasePredictionHead | None, num_classes: int | None, @@ -56,6 +57,9 @@ def __init__( tabular_dim=None, ) + if geo_encoder_prediction: + self.geo_encoder_prediction = geo_encoder_prediction + # Params from alignment model self.match_to_geo = match_to_geo self.ks = ks @@ -72,6 +76,8 @@ def _setup(self, stage: str) -> None: self.geo_encoder.setup() if hasattr(self, "text_encoder"): self.text_encoder.setup() + if hasattr(self, "geo_encoder_prediction"): + self.geo_encoder_prediction.setup() if hasattr(self, "text_encoder") and hasattr(self, "geo_encoder"): # Configure optional extra projection so text embeddings match geo embeddings. @@ -86,9 +92,11 @@ def _setup(self, stage: str) -> None: raise ValueError( "InferenceModel requires `num_classes` to build the prediction head." ) - self.prediction_head.set_dim( - input_dim=self.geo_encoder.output_dim, output_dim=self.num_classes - ) + if hasattr(self, "geo_encoder_prediction"): + input_dim = self.geo_encoder_prediction.output_dim + else: + self.geo_encoder.output_dim + self.prediction_head.set_dim(input_dim=input_dim, output_dim=self.num_classes) self.prediction_head.setup() print("------------------------") @@ -101,21 +109,31 @@ def _step( """Step forward computation of the model.""" pass + def forward_text(self, text: list[str]) -> torch.Tensor: + batch = {"text": text} + return self.text_encoder(batch, "train") + + def forward_geo(self, batched_eo) -> Tuple[torch.Tensor, torch.Tensor | None]: + geo_feats = self.geo_encoder(batched_eo) + pred = None + if hasattr(self, "prediction_head"): + if hasattr(self, "geo_encoder_prediction"): + geo_feats_pr = self.geo_encoder_prediction(batched_eo) + pred = self.prediction_head(geo_feats_pr) + else: + pred = self.prediction_head(geo_feats) + return geo_feats, pred + @override def forward( self, batch: Dict[str, torch.Tensor], - mode: str = "train", - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Model forward logic.""" - - # Embed modalities + ) -> Tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: if hasattr(self, "geo_encoder"): - geo_feats = self.geo_encoder(batch) + geo_feats, pred = self.forward_geo(batch) + if hasattr(self, "text_encoder"): - text_feats = self.text_encoder(batch, mode) - if hasattr(self, "prediction_head"): - pred_feats = self.prediction_head(geo_feats) + text_feats = self.text_encoder(batch, "test") # Change dtype of geo data if it does not match text dtype if ( @@ -125,7 +143,7 @@ def forward( ): geo_feats = geo_feats.to(text_feats.dtype) - return pred_feats, geo_feats, text_feats + return pred, geo_feats, text_feats def concept_similarities(self, geo_embeds, concept=None) -> torch.Tensor: # Get concept embeddings @@ -155,13 +173,15 @@ def _is_prefix_trained(trainable_modules: list[str], prefix: str) -> bool: return any(m.split(".")[0] == prefix for m in trainable_modules) -def load_inference_model(inference_ckpt_path: str) -> InferenceModel: +def load_inference_model(inference_ckpt_path: str, cache_path: str | None) -> InferenceModel: """Loads inference model from a merged checkpoint. :param inference_ckpt_path: path to inference model weights :return: an InferenceModel with pre-trained weights """ inference_ckpt = torch.load(inference_ckpt_path, map_location="cpu", weights_only=False) + if cache_path: + inference_ckpt["hyper_parameters"]["text_encoder"]["hf_cache_dir"] = cache_path model = hydra.utils.instantiate(inference_ckpt["hyper_parameters"]) model.setup("inference") res = model.load_state_dict(inference_ckpt["state_dict"], strict=False) @@ -188,16 +208,20 @@ def merge_inference_model(cfg, save_ckpt=False) -> InferenceModel | None: pred_ckpt = torch.load(pred_ckpt_path, map_location="cpu", weights_only=False) align_ckpt = torch.load(align_ckpt_path, map_location="cpu", weights_only=False) + inference_hparams = {} + geo_encoder_prediction_flag = False # Sanity check: ensure geo encoder configs match. - align_ckpt["hyper_parameters"]["geo_encoder"] = pred_ckpt["hyper_parameters"].get( - "geo_encoder" - ) if pred_ckpt["hyper_parameters"].get("geo_encoder") != align_ckpt["hyper_parameters"].get( "geo_encoder" ): log.warning("Geo encoder configs differ between checkpoints; results may be invalid.") - if input("Do you want to proceed? y/n").lower() == "n": + if input("Do you want to proceed? y/n").lower() != "y": return None + if input("Use separate geo_encoders for prediction and alignment? y/n").lower() == "y": + inference_hparams["geo_encoder_prediction"] = pred_ckpt["hyper_parameters"].get( + "geo_encoder" + ) + geo_encoder_prediction_flag = True pred_trainable_modules = pred_ckpt["hyper_parameters"].get("trainable_modules", []) align_trainable_modules = align_ckpt["hyper_parameters"].get("trainable_modules", []) @@ -209,7 +233,7 @@ def merge_inference_model(cfg, save_ckpt=False) -> InferenceModel | None: raise ValueError("Models are not aligned: both checkpoints trained geo_encoder.") # Instantiate InferenceModel via hydra, using alignment encoder configs with prediction model head configs - inference_hparams = align_ckpt["hyper_parameters"] + inference_hparams.update(align_ckpt["hyper_parameters"]) inference_hparams.update( { "_target_": "src.models.inference_model.InferenceModel", @@ -217,6 +241,7 @@ def merge_inference_model(cfg, save_ckpt=False) -> InferenceModel | None: "num_classes": pred_ckpt["hyper_parameters"].get("num_classes"), } ) + inference_hparams["trainable_modules"] = None inference_hparams["text_encoder"]["hf_cache_dir"] = os.path.join( cfg.paths.cache_dir, "huggingface" ) @@ -228,10 +253,21 @@ def merge_inference_model(cfg, save_ckpt=False) -> InferenceModel | None: text_state = { k: v for k, v in align_ckpt["state_dict"].items() if k.startswith("text_encoder.") } - res = model.load_state_dict(text_state, strict=False) - log_model_loading("text_encoder", res) + if text_state: + res = model.load_state_dict(text_state, strict=False) + log_model_loading("text_encoder", res) - if cfg.training_order[0] == "prediction_model" and not geo_align_encoder_trained: + if geo_encoder_prediction_flag: + geo_state = { + k: v for k, v in align_ckpt["state_dict"].items() if k.startswith("geo_encoder.") + } + geo_pred_state = { + k: v for k, v in pred_ckpt["state_dict"].items() if k.startswith("geo_encoder.") + } + if geo_pred_state: + res = model.load_state_dict(geo_pred_state, strict=False) + log_model_loading("geo_encoder_pred", res) + elif cfg.training_order[0] == "prediction_model" and not geo_align_encoder_trained: geo_state = { k: v for k, v in pred_ckpt["state_dict"].items() if k.startswith("geo_encoder.") } @@ -239,15 +275,17 @@ def merge_inference_model(cfg, save_ckpt=False) -> InferenceModel | None: geo_state = { k: v for k, v in align_ckpt["state_dict"].items() if k.startswith("geo_encoder.") } - res = model.load_state_dict(geo_state, strict=False) - log_model_loading("geo_encoder", res) + if geo_state: + res = model.load_state_dict(geo_state, strict=False) + log_model_loading("geo_encoder", res) # Load prediction head weights from predictive ckpt. head_state = { k: v for k, v in pred_ckpt["state_dict"].items() if k.startswith("prediction_head.") } - res = model.load_state_dict(head_state, strict=False) - log_model_loading("Predictive_head", res) + if head_state: + res = model.load_state_dict(head_state, strict=False) + log_model_loading("Predictive_head", res) # Save model if save_ckpt: diff --git a/src/models/predictive_model.py b/src/models/predictive_model.py index 9e3c28a..a956eb0 100644 --- a/src/models/predictive_model.py +++ b/src/models/predictive_model.py @@ -112,7 +112,9 @@ def _setup_encoders_adapters(self): self.trainable_modules.extend(new_modules) if self.normalize_features: - self.normalizer = nn.LayerNorm(self.geo_encoder.output_dim, dtype=self.geo_encoder.dtype) + self.normalizer = nn.LayerNorm( + self.geo_encoder.output_dim, dtype=self.geo_encoder.dtype + ) self.trainable_modules.append("normalizer") print("Model set up to normalise geo_encoder features.") diff --git a/src/models/text_alignment_model.py b/src/models/text_alignment_model.py index 7a67cde..e693446 100644 --- a/src/models/text_alignment_model.py +++ b/src/models/text_alignment_model.py @@ -1,6 +1,5 @@ from typing import Dict, Tuple, override -import numpy as np import torch import torch.nn.functional as F @@ -11,9 +10,8 @@ RetrievalContrastiveValidation, ) from src.models.components.metrics.metrics_wrapper import MetricsWrapper -from src.models.components.text_encoders.base_text_encoder import ( - BaseTextEncoder, -) +from src.models.components.projectors_adapters.base_encoder import BaseEncoder +from src.models.components.text_encoders.base_text_encoder import BaseTextEncoder class TextAlignmentModel(BaseModel): @@ -26,6 +24,8 @@ def __init__( scheduler: torch.optim.lr_scheduler, loss_fn: BaseLossFn, metrics: MetricsWrapper, + geo_adapter: BaseEncoder | None = None, + text_adapter: BaseEncoder | None = None, num_classes: int | None = None, tabular_dim: int | None = None, ks: list[int] | None = [5, 10, 15], @@ -46,6 +46,7 @@ def __init__( :param match_to_geo: whether to match dimensions of text encoder to geo_encoder or visa- versa """ + super().__init__( trainable_modules=trainable_modules, geo_encoder=geo_encoder, @@ -59,6 +60,8 @@ def __init__( tabular_dim=tabular_dim, ) + self.geo_adapter = geo_adapter + self.text_adapter = text_adapter # Metrics self.ks = ks self.log_kwargs = dict(on_step=False, on_epoch=True, prog_bar=True, sync_dist=True) @@ -76,21 +79,47 @@ def _setup(self, stage: str = "fit") -> None: # Set up encoders and missing adapters/projectors print("-------Model------------") new_modules = [f"geo_encoder.{i}" for i in self.geo_encoder.setup() or []] + + if self.geo_adapter: + self.geo_adapter.set_input_dim(self.geo_encoder.output_dim) + new_modules.extend([f"geo_adapter.{i}" for i in self.geo_adapter.setup() or []]) + new_modules.extend([f"text_encoder.{i}" for i in self.text_encoder.setup() or []]) + if self.text_adapter: + self.text_adapter.set_input_dim(self.text_encoder.input_dim) + new_modules.extend([f"text_adapter.{i}" for i in self.text_adapter.setup() or []]) + self.trainable_modules.extend(new_modules) # Extra projector for text encoder if eo and text dim not match - if self.geo_encoder.output_dim != self.text_encoder.output_dim: - if self.match_to_geo: + geo_branch_dim = ( + self.geo_adapter.output_dim if self.geo_adapter else self.geo_encoder.output_dim + ) + text_branch_dim = ( + self.text_adapter.output_dim if self.text_adapter else self.text_encoder.output_dim + ) + + if geo_branch_dim != text_branch_dim: + if self.geo_adapter or self.text_adapter: + print( + f"You opted to use:{' geo' if self.geo_adapter else '' and ' text' if self.text_adapter else ''} adapter", + "but you miss-configured output dimensions:\n" + f"geo: {geo_branch_dim} vs text: {text_branch_dim}\n", + "Please try again.", + ) + elif self.match_to_geo: self.text_encoder.add_projector(projected_dim=self.geo_encoder.output_dim) self.trainable_modules.append("text_encoder.extra_projector") else: self.geo_encoder.add_projector(projected_dim=self.text_encoder.output_dim) self.trainable_modules.append("geo_encoder.extra_projector") + print("------------------------") + + def on_fit_start(self): # Configure contrastive retrieval evaluation self.setup_retrieval_evaluation(verbose=0) - print("------------------------") + print("Retrieval evaluation configured") def setup_retrieval_evaluation( self, @@ -124,8 +153,10 @@ def setup_retrieval_evaluation( return # Encode concepts if text branch is frozen - with torch.no_grad(): + with torch.inference_mode(): self.concept_embeds = self.text_encoder({"text": self.concepts}, mode="train") + self.concept_embeds = F.normalize(self.concept_embeds, dim=1) + self.concept_embeds = self.concept_embeds.to(self.device) @override def forward( @@ -137,7 +168,11 @@ def forward( # Embed modalities geo_feats = self.geo_encoder(batch) + if self.geo_adapter: + geo_feats = self.geo_adapter(geo_feats) text_feats = self.text_encoder(batch, mode) + if self.text_adapter: + text_feats = self.text_adapter(text_feats) # Change dtype of geo data if it does not match text dtype if geo_feats.dtype != text_feats.dtype: @@ -160,7 +195,21 @@ def _step(self, batch: Dict[str, torch.Tensor], mode: str = "train"): geo_feats, text_feats = feats[0], feats[1] # Get loss - loss = self.loss_fn(geo_feats, text_feats) + aux_values = batch["aux"].get("aux") + aux_ids_per_caption = batch.get("text_aux_ids") + + if geo_feats.isnan().any(): + print(geo_feats) + print(batch["name_loc"]) + exit() + + loss = self.loss_fn( + geo_feats, + text_feats, + mode=mode, + aux_values=aux_values, + aux_ids_per_caption=aux_ids_per_caption, + ) # Get similarities with torch.no_grad(): @@ -185,10 +234,12 @@ def _step(self, batch: Dict[str, torch.Tensor], mode: str = "train"): self.log_dict(metrics, batch_size=local_batch_size, **self.log_kwargs) if mode in ["val", "test"]: + aux = batch.get("aux", {}).get("aux") self.outputs_epoch_memory.append( { - "geo_feats": geo_feats.detach(), - "aux_vals": batch.get("aux", {}).get("aux").detach(), + # Store on CPU to avoid holding the whole epoch on GPU. + "geo_feats": geo_feats.detach().cpu(), + "aux_vals": aux.detach().cpu() if aux is not None else None, } ) @@ -198,8 +249,11 @@ def _on_epoch_end(self, mode: str, verbose=0): # Combine batches geo_feats = torch.cat([x["geo_feats"] for x in self.outputs_epoch_memory], dim=0) + geo_feats = geo_feats.to(self.device, non_blocking=True) - aux_vals = torch.cat([x["aux_vals"] for x in self.outputs_epoch_memory], dim=0) + aux_vals = torch.cat([x["aux_vals"] for x in self.outputs_epoch_memory], dim=0).to( + self.device, non_blocking=True + ) # Rank on similarity similarity = self.concept_similarities(geo_feats) @@ -247,23 +301,34 @@ def on_test_epoch_end(self): return self._on_epoch_end("test") def concept_similarities(self, geo_embeds, concept=None) -> torch.Tensor: + device_type = geo_embeds.device.type + is_bf16 = self.trainer.precision == "bf16-mixed" and device_type == "cuda" + # Get concept embeddings if concept is not None: # If only one concept is provided if isinstance(concept, str): concept = [concept] - with torch.no_grad(): - concept_embeds = self.text_encoder({"text": concept}, mode="train") + + with torch.inference_mode(): + with torch.autocast( + device_type=device_type, dtype=torch.bfloat16, enabled=is_bf16 + ): + concept_embeds = self.text_encoder({"text": concept}, mode="train") + concept_embeds = F.normalize(concept_embeds, dim=1) elif self.concept_embeds is not None: concept_embeds = self.concept_embeds else: - with torch.no_grad(): - concept_embeds = self.text_encoder({"text": self.concepts}, mode="train") + with torch.inference_mode(): + with torch.autocast( + device_type=device_type, dtype=torch.bfloat16, enabled=is_bf16 + ): + concept_embeds = self.text_encoder({"text": self.concepts}, mode="train") + concept_embeds = F.normalize(concept_embeds, dim=1) # Similarity geo_embeds = F.normalize(geo_embeds, dim=1) - concept_embeds = F.normalize(concept_embeds, dim=1) similarity_matrix = concept_embeds @ geo_embeds.T return similarity_matrix diff --git a/src/train.py b/src/train.py index 949e21b..f29af2f 100644 --- a/src/train.py +++ b/src/train.py @@ -10,7 +10,7 @@ from omegaconf import DictConfig, OmegaConf from src.data.base_datamodule import BaseDataModule -from src.utils.experiment_tracking import experiment_check +from src.utils.experiment_tracking import compose_experiment_name, experiment_check rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) load_dotenv() @@ -96,6 +96,8 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: log.info("Logging hyperparameters!") log_hyperparameters(object_dict) group = cfg.get("experiment_name", "null") + if group == "null": + compose_experiment_name(cfg) wandb_logger.log_metrics({"experiment": group}) if cfg.get("train"): @@ -118,10 +120,14 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: else: data_name = "satclip" else: - data_name = f"{k}_{data_dict[k]['size']}" + data_name = f"{k}_{data_dict[k].get('size', '')}" else: ks = list(data_dict.keys()) - data_name = str("".join([f'{k}_{data_dict[k].get("size", "")}' for k in ks])) + ks_new = [ + f"{k}{f'_{k.get('size')}' if isinstance(k, dict) and k.get('size') else ''}" + for k in ks + ] + data_name = "-".join(map(str, ks_new)) # Log details to wandb wandb_logger.log_metrics( @@ -135,7 +141,25 @@ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: train_metrics = trainer.callback_metrics - if cfg.get("test"): + if cfg.get("validate") and wandb_logger is not None: + # Run validation with the best ckpt + log.info("Validating the best ckpt!") + ckpt_path = trainer.checkpoint_callback.best_model_path + if ckpt_path == "": + log.warning("Best ckpt not found! Using current weights for testing...") + ckpt_path = None + + trainer.validate( + model=model, + datamodule=datamodule, + ckpt_path=ckpt_path, + weights_only=False, + ) + + val_metrics = trainer.callback_metrics + wandb_logger.log_metrics({f"best_val_{k}": v for k, v in val_metrics.items()}) + + if cfg.get("test") and wandb_logger is not None: log.info("Starting testing!") ckpt_path = trainer.checkpoint_callback.best_model_path if ckpt_path == "": diff --git a/src/utils/experiment_tracking.py b/src/utils/experiment_tracking.py index a19ccc2..4f99ae0 100644 --- a/src/utils/experiment_tracking.py +++ b/src/utils/experiment_tracking.py @@ -1,6 +1,5 @@ import os import re -import shutil from pathlib import Path import pandas as pd @@ -13,6 +12,92 @@ ENTITY = "aether_xai" +def parse_data_name(run=None, cfg=None): + if cfg is None: + assert run is not None + data_dict = run.config["data"]["dataset"]["modalities"] + else: + data_dict = cfg["data"]["dataset"]["modalities"] + + if len(data_dict) == 1: + k = list(data_dict.keys())[0] + if k == "coords": + if "GeoClip" in run.config["model"]["geo_encoder"]["_target_"]: + data_name = "geoclip" + else: + data_name = "satclip" + else: + ks = list(data_dict.keys()) + ks_new = [ + f"{k}{f'_{k.get('size')}' if isinstance(k, dict) and k.get('size') else ''}" + for k in ks + ] + data_name = "-".join(map(str, ks_new)) + + return data_name + + +def compose_experiment_name(cfg: DictConfig) -> str: + """For alignment + (unlab_)(geo_encoder)_(eo_mod)_(geo_adapter)_(text_encoder)_(text_adapter)_(batch_size) # noqa: + + RST306. + """ + + name = "" + + # Unlabeled + if cfg.data.dataset.use_unlabelled_data: + name += "unlab_" # noqa: RST306 + + # geo encoder + name += cfg.model.geo_encoder._target_.split(".")[-1] + + # eo mod + name += "_" + parse_data_name(cfg=cfg) + + # batch size + if cfg.model._target_ == "src.models.text_alignment_model.TextAlignmentModel": + # geo adapter + if cfg.model.get("geo_adapter") is not None: + name += "_" + cfg.model.geo_adapter._target_.split(".")[-1] + + # text encoder + if cfg.model.get("text_adapter") is not None: + name += "_" + cfg.model.text_encoder._target_.split(".")[-1] + + # text adapter + if cfg.model.get("text_adapter") is not None: + name += "_" + cfg.model.text_adapter._target_.split(".")[-1] + + # batch size + name += f"_b-{cfg.data.batch_size}" + else: + # prediction head + if cfg.model.get("prediction_head") is not None: + name += "_" + cfg.model.prediction_head._target_.split(".")[-1] + + # extras for experiment identification + if cfg.get("experiment_name_extra"): + name += "_" + cfg.experiment_name_extra + + return name + + +def update_wandb_run_id(run_id, project_name, updates_dict): + api = wandb.Api() + entity = ENTITY + + run = api.run(f"{entity}/{project_name}/{run_id}") + + # Update the config + for key, value in updates_dict.items(): + run.config[key] = value + + run.config.save() + print(f"Successfully updated run {run_id}") + + def experiment_check(cfg: DictConfig): """Check if requested through configs experiments has already been run.""" # Get experiments from wandb (cross-device sync) @@ -40,6 +125,7 @@ def get_experiments_from_wandb(cfg: DictConfig) -> pd.DataFrame | None: # Get existing experiment table df, df_path = open_experiment_df(cfg) ids = list(df.run_id) if len(df) > 0 else [] + df_len = len(df) # Connect to wandb api api = wandb.Api() @@ -79,9 +165,10 @@ def get_experiments_from_wandb(cfg: DictConfig) -> pd.DataFrame | None: if best_val_loss is None: if len(results) > 0: loss_per_epoch = results.groupby("epoch")["val_loss"].mean() - if best_epoch != loss_per_epoch.idxmin(): - raise ValueError() - best_val_loss = loss_per_epoch.min().item() + if best_epoch: + best_val_loss = loss_per_epoch[best_epoch] + else: + best_val_loss = loss_per_epoch.min().item() run.summary.update({"best_val_loss": best_val_loss}) if project == "s2bms_prediction" and best_val_mse_loss is None: mse_loss_per_epoch = results.groupby("epoch")["val_mse_loss"].mean() @@ -92,17 +179,7 @@ def get_experiments_from_wandb(cfg: DictConfig) -> pd.DataFrame | None: if experiment is None: experiment = run.config["experiment_name"] run.summary.update({"experiment": experiment}) - - data_dict = run.config["data"]["dataset"]["modalities"] - if len(data_dict) == 1: - k = list(data_dict.keys())[0] - if k == "coords": - if "GeoClip" in run.config["model"]["geo_encoder"]["_target_"]: - data_name = "geoclip" - else: - data_name = "satclip" - else: - data_name = f"{k}_{data_dict[k]['size']}" + data_name = parse_data_name(run=run) run.summary.update({"data_used": data_name}) # Add missing experiments to the table @@ -119,9 +196,12 @@ def get_experiments_from_wandb(cfg: DictConfig) -> pd.DataFrame | None: ) # Return experiment df - if len(df) > 0: + if df_len - len(df) > 0: print(f"Saved {len(runs_list)} runs to {df_path}.") return df + if len(df) > 0: + print("No runs updated.") + return df else: print("No runs found.") return None @@ -209,6 +289,6 @@ def clean_local_ckpts(cfg: DictConfig, df: pd.DataFrame) -> None: if str(local_ckpt) not in keep: print(f"Do you want to remove {local_ckpt}? (y/n)") answer = input() - if answer != "y": - shutil.rmtree(local_ckpt) + if answer == "y": + os.remove(str(local_ckpt)) print(f"Removed {local_ckpt}.") diff --git a/src/utils/logging_utils.py b/src/utils/logging_utils.py index f23f4df..18efa73 100644 --- a/src/utils/logging_utils.py +++ b/src/utils/logging_utils.py @@ -70,16 +70,25 @@ def _group_keys(keys: list[str]) -> dict[str, list[str]]: return grouped -def log_model_loading(tag: str, result) -> None: +def log_model_loading(tag: str, result, printing=False) -> None: """Log missing/unexpected keys from `load_state_dict`.""" missing_keys, unexpected_keys = result if missing_keys: grouped = _group_keys(list(missing_keys)) summary = {k: len(v) for k, v in grouped.items()} - log.warning(f"[{tag}] Missing keys: {summary}") + if printing: + print(f"[{tag}] Missing keys: {summary}") + else: + log.warning(f"[{tag}] Missing keys: {summary}") if unexpected_keys: grouped = _group_keys(list(unexpected_keys)) summary = {k: len(v) for k, v in grouped.items()} - log.warning(f"[{tag}] Unexpected keys: {summary}") + if printing: + print(f"[{tag}] Unexpected keys: {summary}") + else: + log.warning(f"[{tag}] Unexpected keys: {summary}") if not missing_keys and not unexpected_keys: - log.info(f"[{tag}] Module weights loaded successfully") + if printing: + print(f"[{tag}] Module weights loaded successfully") + else: + log.info(f"[{tag}] Module weights loaded successfully") diff --git a/src/utils/utils.py b/src/utils/utils.py index f1771a2..a34d208 100644 --- a/src/utils/utils.py +++ b/src/utils/utils.py @@ -119,8 +119,9 @@ def wrap(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]: if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.ipc_collect() - except Exception: + except Exception as e: # cleanup should never mask the original exception + print(e) pass return metric_dict, object_dict diff --git a/tests/test_dynamic_gate_fusion_encoder.py b/tests/test_dynamic_gate_fusion_encoder.py index 6dff786..30e4893 100644 --- a/tests/test_dynamic_gate_fusion_encoder.py +++ b/tests/test_dynamic_gate_fusion_encoder.py @@ -17,6 +17,7 @@ from src.models.components.geo_encoders.encoder_wrapper import EncoderWrapper from src.models.components.geo_encoders.tabular_encoder import TabularEncoder + # --------------------------------------------------------------------------- # Minimal stub encoder for testing — outputs a fixed-size embedding. # --------------------------------------------------------------------------- @@ -169,7 +170,7 @@ def test_with_tabular_encoder(): Both branches projected to the same output_dim (32) via per-branch projectors. """ - from src.models.components.geo_encoders.mlp_projector import MLPProjector + from src.models.components.projectors_adapters.mlp_projector import MLPProjector tabular_dim = 23 tab_enc = TabularEncoder(output_dim=32) diff --git a/tests/test_gated_fusion_encoder.py b/tests/test_gated_fusion_encoder.py index 9fc2169..4593d0d 100644 --- a/tests/test_gated_fusion_encoder.py +++ b/tests/test_gated_fusion_encoder.py @@ -15,6 +15,7 @@ from src.models.components.geo_encoders.encoder_wrapper import EncoderWrapper from src.models.components.geo_encoders.tabular_encoder import TabularEncoder + # --------------------------------------------------------------------------- # Minimal stub encoder for testing — outputs a fixed-size embedding. # --------------------------------------------------------------------------- @@ -137,7 +138,7 @@ def test_with_tabular_encoder(): Both branches projected to the same output_dim (32) via per-branch projectors. """ - from src.models.components.geo_encoders.mlp_projector import MLPProjector + from src.models.components.projectors_adapters.mlp_projector import MLPProjector tabular_dim = 23 tab_enc = TabularEncoder(output_dim=32) diff --git a/tests/test_geo_encoders.py b/tests/test_geo_encoders.py index 4f062fb..5490eac 100644 --- a/tests/test_geo_encoders.py +++ b/tests/test_geo_encoders.py @@ -3,7 +3,6 @@ from src.models.components.geo_encoders.average_encoder import AverageEncoder from src.models.components.geo_encoders.cnn_encoder import CNNEncoder from src.models.components.geo_encoders.geoclip import GeoClipCoordinateEncoder -from src.models.components.geo_encoders.mlp_projector import MLPProjector from src.models.components.geo_encoders.tabular_encoder import TabularEncoder @@ -15,16 +14,13 @@ def test_geo_encoder_generic_properties(create_butterfly_dataset): "cnn": CNNEncoder, "average": AverageEncoder, "tabular": TabularEncoder, - "mlp_projector": MLPProjector, } ds, dm = create_butterfly_dataset dm.setup() batch = next(iter(dm.train_dataloader())) for geo_encoder_name, geo_encoder_class in dict_geo_encoders.items(): - if geo_encoder_class is MLPProjector: - geo_encoder = geo_encoder_class(output_dim=64, input_dim=128) - elif geo_encoder_class is TabularEncoder: + if geo_encoder_class is TabularEncoder: geo_encoder = geo_encoder_class(output_dim=64, input_dim=128, hidden_dim=128) else: geo_encoder = geo_encoder_class()