Rewrite: GatedFusion architecture + UAV-GeoLoc dataset

Architecture v2:
- Query branch: drone + text -> GatedFusion -> proj -> query_emb
- Gallery branch: satellite -> proj -> gallery_emb
- Single InfoNCE loss (asymmetric 0.6/0.4)
- GatedFusion: learnable gated addition (sigma(alpha)*img + (1-sigma(alpha))*text)
- Baseline mode: gate=1.0 (text ignored)

Dataset:
- UAV-GeoLoc loader with template captions from path metadata
- 27 terrain types with predefined features
- Random positive crop sampling per epoch

Configs: balanced (gate=0.7), baseline (no text), text_heavy (gate=0.3)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-17 17:13:00 +03:00
parent 2ce4017ebd
commit abb3337f8d
12 changed files with 1077 additions and 781 deletions

View File

@@ -1,5 +1,5 @@
# Balanced configuration — primary test setup.
# L = 1.0 * L_img_img + 0.3 * L_sat_cap + 0.3 * L_drone_cap + 0.1 * L_cap_cap
# Balanced: GatedFusion with text captions enabled.
# query = sigma(alpha) * drone + (1-sigma(alpha)) * text -> InfoNCE vs satellite
import src.datasets.visloc_with_captions
import src.losses.multi_infonce
@@ -12,43 +12,37 @@ DualEncoderCaptionTest.pretrained_path = "checkpoints/RS5M_ViT-B-32.pt"
DualEncoderCaptionTest.unfreeze_mode = "last_block"
DualEncoderCaptionTest.embed_dim = 512
DualEncoderCaptionTest.use_mlp_heads = False
DualEncoderCaptionTest.shared_image_head = True
DualEncoderCaptionTest.baseline_mode = False
DualEncoderCaptionTest.init_gate = 0.7
DualEncoderCaptionTest.device = "cuda"
ProjectionHead.in_dim = 512
ProjectionHead.out_dim = 512
ProjectionHead.use_mlp = False
# ---- Fusion ----
GatedFusion.init_gate = 0.7
GatedFusion.baseline_mode = False
# ---- Loss ----
MultiTermInfoNCE.temperature_init = 0.1
MultiTermInfoNCE.temperature_final = 0.01
MultiTermInfoNCE.label_smoothing = 0.1
MultiTermInfoNCE.asym_drone_to_sat = 0.6
MultiTermInfoNCE.asym_sat_to_drone = 0.4
MultiTermInfoNCE.warmup_epochs = 3
MultiTermInfoNCE.text_ramp_epochs = 10
MultiTermInfoNCE.lambda_ii = 1.0
MultiTermInfoNCE.lambda_sc_max = 0.3
MultiTermInfoNCE.lambda_dc_max = 0.3
MultiTermInfoNCE.lambda_cc_max = 0.1
InfoNCELoss.temperature_init = 0.1
InfoNCELoss.temperature_final = 0.01
InfoNCELoss.label_smoothing = 0.1
InfoNCELoss.weight_q2g = 0.6
InfoNCELoss.weight_g2q = 0.4
# ---- Dataset ----
VisLocCaptionDataset.caption_strategy = "hybrid"
VisLocCaptionDataset.drop_caption_prob = 0.0
VisLocCaptionDataset.seed = 42
GeoLocCaptionDataset.drop_caption_prob = 0.0
GeoLocCaptionDataset.seed = 42
# ---- Training ----
TrainConfig.train_manifest = "data/visloc_train.json"
TrainConfig.val_manifest = "data/visloc_val.json"
TrainConfig.image_root = "data/visloc/images"
TrainConfig.train_query_file = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/Index/train_query.txt"
TrainConfig.val_query_file = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc/Index/val_query.txt"
TrainConfig.data_root = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc"
TrainConfig.output_dir = "out/caption_test/balanced"
TrainConfig.epochs = 30
TrainConfig.epochs = 10
TrainConfig.batch_size = 128
TrainConfig.num_workers = 4
TrainConfig.learning_rate = 1e-4
TrainConfig.weight_decay = 1e-4
TrainConfig.grad_clip = 1.0
TrainConfig.use_amp = True
TrainConfig.eval_every = 1
TrainConfig.eval_every = 2
TrainConfig.seed = 42
TrainConfig.device = "cuda"