abb3337f8d
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>
49 lines
1.6 KiB
Plaintext
49 lines
1.6 KiB
Plaintext
# 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
|
|
import src.models.dual_encoder
|
|
import src.training.train
|
|
|
|
# ---- Dual encoder ----
|
|
DualEncoderCaptionTest.variant = "ViT-B-32"
|
|
DualEncoderCaptionTest.pretrained_path = "checkpoints/RS5M_ViT-B-32.pt"
|
|
DualEncoderCaptionTest.unfreeze_mode = "last_block"
|
|
DualEncoderCaptionTest.embed_dim = 512
|
|
DualEncoderCaptionTest.use_mlp_heads = False
|
|
DualEncoderCaptionTest.baseline_mode = False
|
|
DualEncoderCaptionTest.init_gate = 0.7
|
|
DualEncoderCaptionTest.device = "cuda"
|
|
|
|
# ---- Fusion ----
|
|
GatedFusion.init_gate = 0.7
|
|
GatedFusion.baseline_mode = False
|
|
|
|
# ---- Loss ----
|
|
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 ----
|
|
GeoLocCaptionDataset.drop_caption_prob = 0.0
|
|
GeoLocCaptionDataset.seed = 42
|
|
|
|
# ---- Training ----
|
|
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 = 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 = 2
|
|
TrainConfig.seed = 42
|
|
TrainConfig.device = "cuda"
|