Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rewrite: GatedFusion architecture + UAV-GeoLoc dataset

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

208
src/losses/multi_infonce.py
View File

@@ -1,14 +1,10 @@
from __future__ import annotations
"""Multi-term InfoNCE loss for caption quality validation.
"""InfoNCE loss for cross-view geo-localization with optional text fusion.
Four InfoNCE terms over projected embeddings:
L = lambda_ii * L_img_img
+ lambda_sc * L_sat_cap
+ lambda_dc * L_drone_cap
+ lambda_cc * L_cap_cap
where L_img_img is the classical symmetric CVGL contrastive loss
with asymmetric weights (0.6 drone->sat + 0.4 sat->drone).
Single symmetric InfoNCE between query (drone+text fused) and gallery (satellite).
Asymmetric weighting: query->gallery weighted higher (real use-case direction).
Cosine temperature schedule for sharper distribution over training.
"""
import math
@@ -27,26 +23,12 @@ def _symmetric_info_nce(
weight_a2b: float = 0.5,
weight_b2a: float = 0.5,
) -> torch.Tensor:
"""Compute weighted symmetric InfoNCE between two L2-normalized embeddings.
Args:
emb_a: First embedding set [B, D].
emb_b: Second embedding set [B, D]. Positive pairs are on the diagonal.
temperature: Softmax temperature (smaller = sharper distribution).
label_smoothing: Cross-entropy label smoothing epsilon.
weight_a2b: Weight for A-query direction.
weight_b2a: Weight for B-query direction.
Returns:
Scalar weighted loss.
"""
"""Weighted symmetric InfoNCE. Positives on the diagonal."""
batch_size = emb_a.size(0)
logits = emb_a @ emb_b.t() / temperature
targets = torch.arange(batch_size, device=emb_a.device)
loss_a2b = F.cross_entropy(logits, targets, label_smoothing=label_smoothing)
loss_b2a = F.cross_entropy(logits.t(), targets, label_smoothing=label_smoothing)
return weight_a2b * loss_a2b + weight_b2a * loss_b2a
@@ -56,85 +38,23 @@ def cosine_temperature(
tau_init: float = 0.1,
tau_final: float = 0.01,
) -> float:
"""Cosine-decay schedule for InfoNCE temperature.
Args:
epoch: Current training epoch (0-indexed).
total_epochs: Total number of epochs.
tau_init: Initial temperature.
tau_final: Final temperature.
Returns:
Temperature value for this epoch.
"""
"""Cosine-decay schedule for InfoNCE temperature."""
total_epochs = max(total_epochs, 1)
progress = min(max(epoch / total_epochs, 0.0), 1.0)
cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
return tau_final + (tau_init - tau_final) * cosine
def curriculum_lambdas(
epoch: int,
warmup_epochs: int = 3,
text_ramp_epochs: int = 10,
lambda_ii: float = 1.0,
lambda_sc_max: float = 0.3,
lambda_dc_max: float = 0.3,
lambda_cc_max: float = 0.1,
) -> dict[str, float]:
"""Compute per-epoch loss weights under the curriculum schedule.
- Epochs 0..warmup_epochs: image-image only.
- Epochs warmup..text_ramp_epochs: linearly ramp sat-cap and drone-cap.
- Epochs >= text_ramp_epochs: full loss including caption-caption term.
Args:
epoch: Current epoch (0-indexed).
warmup_epochs: Number of warmup epochs (no text losses).
text_ramp_epochs: Epoch when text losses reach max.
lambda_ii: Constant weight for image-image loss.
lambda_sc_max: Max weight for satellite-caption loss.
lambda_dc_max: Max weight for drone-caption loss.
lambda_cc_max: Max weight for caption-caption loss.
Returns:
Dict with keys 'img_img', 'sat_cap', 'drone_cap', 'cap_cap'.
"""
if epoch < warmup_epochs:
ramp = 0.0
elif epoch >= text_ramp_epochs:
ramp = 1.0
else:
denom = max(text_ramp_epochs - warmup_epochs, 1)
ramp = (epoch - warmup_epochs) / denom
return {
"img_img": lambda_ii,
"sat_cap": lambda_sc_max * ramp,
"drone_cap": lambda_dc_max * ramp,
"cap_cap": lambda_cc_max * ramp,
}
@gin.configurable
class MultiTermInfoNCE(nn.Module):
"""Multi-term InfoNCE loss with curriculum and cosine temperature.
Produces total loss and per-component diagnostics. All inputs must be
L2-normalized embeddings of the same dimension.
class InfoNCELoss(nn.Module):
"""Symmetric InfoNCE with cosine temperature schedule.
Args:
temperature_init: Initial temperature (epoch 0).
temperature_final: Final temperature after cosine decay.
label_smoothing: Cross-entropy label smoothing epsilon.
asym_drone_to_sat: Weight for drone->sat InfoNCE direction.
asym_sat_to_drone: Weight for sat->drone InfoNCE direction.
warmup_epochs: Epochs with image-image loss only.
text_ramp_epochs: Epoch at which text losses reach max.
lambda_ii: Constant weight for image-image loss.
lambda_sc_max: Max weight for sat-caption loss.
lambda_dc_max: Max weight for drone-caption loss.
lambda_cc_max: Max weight for caption-caption loss.
temperature_init: Temperature at epoch 0.
temperature_final: Temperature after cosine decay.
label_smoothing: Cross-entropy label smoothing.
weight_q2g: Weight for query->gallery direction.
weight_g2q: Weight for gallery->query direction.
"""
def __init__(
@@ -142,27 +62,15 @@ class MultiTermInfoNCE(nn.Module):
temperature_init: float = 0.1,
temperature_final: float = 0.01,
label_smoothing: float = 0.1,
asym_drone_to_sat: float = 0.6,
asym_sat_to_drone: float = 0.4,
warmup_epochs: int = 3,
text_ramp_epochs: int = 10,
lambda_ii: float = 1.0,
lambda_sc_max: float = 0.3,
lambda_dc_max: float = 0.3,
lambda_cc_max: float = 0.1,
weight_q2g: float = 0.6,
weight_g2q: float = 0.4,
) -> None:
super().__init__()
self.temperature_init = temperature_init
self.temperature_final = temperature_final
self.label_smoothing = label_smoothing
self.asym_drone_to_sat = asym_drone_to_sat
self.asym_sat_to_drone = asym_sat_to_drone
self.warmup_epochs = warmup_epochs
self.text_ramp_epochs = text_ramp_epochs
self.lambda_ii = lambda_ii
self.lambda_sc_max = lambda_sc_max
self.lambda_dc_max = lambda_dc_max
self.lambda_cc_max = lambda_cc_max
self.weight_q2g = weight_q2g
self.weight_g2q = weight_g2q
def forward(
self,
@@ -170,17 +78,16 @@ class MultiTermInfoNCE(nn.Module):
epoch: int,
total_epochs: int,
) -> dict[str, torch.Tensor]:
"""Compute multi-term loss.
"""Compute InfoNCE loss.
Args:
embeddings: Dict with keys 'drone', 'sat', and optionally
'cap_drone', 'cap_sat'. Each [B, D] L2-normalized.
embeddings: Dict with 'query' and 'gallery' [B, D] L2-normalized,
plus 'gate' (float) from fusion module.
epoch: Current epoch (0-indexed).
total_epochs: Total epochs for temperature schedule.
Returns:
Dict with scalar tensors: 'total', 'img_img', 'sat_cap',
'drone_cap', 'cap_cap', plus 'temperature' and 'lambdas'.
Dict with 'total', 'temperature', 'gate'.
"""
tau = cosine_temperature(
epoch=epoch,
@@ -188,75 +95,20 @@ class MultiTermInfoNCE(nn.Module):
tau_init=self.temperature_init,
tau_final=self.temperature_final,
)
lambdas = curriculum_lambdas(
epoch=epoch,
warmup_epochs=self.warmup_epochs,
text_ramp_epochs=self.text_ramp_epochs,
lambda_ii=self.lambda_ii,
lambda_sc_max=self.lambda_sc_max,
lambda_dc_max=self.lambda_dc_max,
lambda_cc_max=self.lambda_cc_max,
)
drone = embeddings["drone"]
sat = embeddings["sat"]
# Image-image symmetric InfoNCE with asymmetric weights.
loss_ii = _symmetric_info_nce(
emb_a=drone,
emb_b=sat,
loss = _symmetric_info_nce(
emb_a=embeddings["query"],
emb_b=embeddings["gallery"],
temperature=tau,
label_smoothing=self.label_smoothing,
weight_a2b=self.asym_drone_to_sat,
weight_b2a=self.asym_sat_to_drone,
weight_a2b=self.weight_q2g,
weight_b2a=self.weight_g2q,
)
loss_sc = torch.zeros_like(loss_ii)
loss_dc = torch.zeros_like(loss_ii)
loss_cc = torch.zeros_like(loss_ii)
if "cap_sat" in embeddings and lambdas["sat_cap"] > 0:
loss_sc = _symmetric_info_nce(
emb_a=sat,
emb_b=embeddings["cap_sat"],
temperature=tau,
label_smoothing=self.label_smoothing,
)
if "cap_drone" in embeddings and lambdas["drone_cap"] > 0:
loss_dc = _symmetric_info_nce(
emb_a=drone,
emb_b=embeddings["cap_drone"],
temperature=tau,
label_smoothing=self.label_smoothing,
)
if (
"cap_drone" in embeddings
and "cap_sat" in embeddings
and lambdas["cap_cap"] > 0
):
loss_cc = _symmetric_info_nce(
emb_a=embeddings["cap_drone"],
emb_b=embeddings["cap_sat"],
temperature=tau,
label_smoothing=self.label_smoothing,
)
total = (
lambdas["img_img"] * loss_ii
+ lambdas["sat_cap"] * loss_sc
+ lambdas["drone_cap"] * loss_dc
+ lambdas["cap_cap"] * loss_cc
)
gate = embeddings.get("gate", 1.0)
return {
"total": total,
"img_img": loss_ii.detach(),
"sat_cap": loss_sc.detach(),
"drone_cap": loss_dc.detach(),
"cap_cap": loss_cc.detach(),
"temperature": torch.tensor(tau, device=total.device),
"lambda_ii": torch.tensor(lambdas["img_img"], device=total.device),
"lambda_sc": torch.tensor(lambdas["sat_cap"], device=total.device),
"lambda_dc": torch.tensor(lambdas["drone_cap"], device=total.device),
"lambda_cc": torch.tensor(lambdas["cap_cap"], device=total.device),
"total": loss,
"temperature": torch.tensor(tau, device=loss.device),
"gate": torch.tensor(gate, device=loss.device),
}