From 04d53072212f5a262994d9bd0580a1d46bc46d72 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Fri, 24 Apr 2026 08:47:33 +0300 Subject: [PATCH] Asymmetric encoders + MONA all 24 blocks + 1024-dim + hard negatives Architecture changes: - Asymmetric DINOv3: WEB (drone) + SAT (satellite) with separate MONA - MONA on all 24 blocks per encoder (was last 12) - Remove projection, native 1024-dim retrieval space (was 512) - Total: 748M params, 17.6M trainable (2.35%) Hard negative memory bank: - MoCo-style FIFO queue of 4096 detached gallery embeddings - Each batch: B in-batch + Q queue negatives in InfoNCE - Queue updated after each forward pass Training config: - batch_size=8, grad_accum=8, effective_batch=64 - eval_every=1 (eval + train recall every epoch) - Max bs=24 with grad checkpointing on RTX 4090 Co-Authored-By: Claude Opus 4.6 (1M context) --- conf/gtauav_balanced.gin | 11 ++--- src/losses/hard_negatives.py | 70 ++++++++++++++++++++++++++++++++ src/losses/multi_infonce.py | 44 +++++++++++++++----- src/models/asymmetric_encoder.py | 39 +++++++----------- src/training/train_gtauav.py | 19 ++++++++- 5 files changed, 142 insertions(+), 41 deletions(-) create mode 100644 src/losses/hard_negatives.py diff --git a/conf/gtauav_balanced.gin b/conf/gtauav_balanced.gin index 39218c4..a958a2c 100644 --- a/conf/gtauav_balanced.gin +++ b/conf/gtauav_balanced.gin @@ -1,6 +1,6 @@ -# GTA-UAV Balanced: GatedFusion with L1/L2/L3 captions on both branches. +# GTA-UAV Balanced: Asymmetric DINOv3 (WEB+SAT) with L1/L2/L3 captions. # query = sigma(alpha) * drone + (1-sigma(alpha)) * text -> InfoNCE vs gallery -# 10 epochs, DINOv3 + DGTRS-CLIP, MONA + LoRA adapters. +# 10 epochs, MONA all 24 blocks, 1024-dim retrieval, hard negative bank. # # NOTE: TrainConfigGTAUAV is registered by train_gtauav.py before gin parsing. # InfoNCELoss is registered via import below. @@ -15,9 +15,9 @@ TrainConfigGTAUAV.learning_rate = 1e-4 TrainConfigGTAUAV.text_lr_factor = 0.1 TrainConfigGTAUAV.weight_decay = 1e-4 TrainConfigGTAUAV.grad_clip = 1.0 -TrainConfigGTAUAV.grad_accum_steps = 1 +TrainConfigGTAUAV.grad_accum_steps = 8 TrainConfigGTAUAV.use_amp = True -TrainConfigGTAUAV.eval_every = 2 +TrainConfigGTAUAV.eval_every = 1 TrainConfigGTAUAV.warmup_epochs = 2 TrainConfigGTAUAV.seed = 42 TrainConfigGTAUAV.device = "cuda" @@ -25,7 +25,7 @@ TrainConfigGTAUAV.device = "cuda" # ---- Model ---- TrainConfigGTAUAV.init_gate = 0.7 TrainConfigGTAUAV.baseline_mode = False -TrainConfigGTAUAV.shared_encoder = True +TrainConfigGTAUAV.shared_encoder = False TrainConfigGTAUAV.gradient_checkpointing = True # ---- Loss ---- @@ -34,6 +34,7 @@ TrainConfigGTAUAV.label_smoothing = 0.1 TrainConfigGTAUAV.weight_q2g = 0.6 TrainConfigGTAUAV.weight_g2q = 0.4 TrainConfigGTAUAV.learnable_temperature = True +TrainConfigGTAUAV.neg_bank_size = 4096 # ---- Output ---- TrainConfigGTAUAV.output_dir = "out/gtauav/with_text" diff --git a/src/losses/hard_negatives.py b/src/losses/hard_negatives.py new file mode 100644 index 0000000..18646ae --- /dev/null +++ b/src/losses/hard_negatives.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +"""Hard negative memory bank for contrastive learning. + +MoCo-style FIFO queue of recent gallery embeddings. Each batch gets +B in-batch negatives + Q queue negatives, significantly increasing +the effective number of negatives without extra VRAM for forward pass. + +Usage: + bank = NegativeMemoryBank(size=4096, dim=1024) + # In training loop: + sim = bank.compute_similarity(query, gallery) # [B, B + Q] + bank.enqueue(gallery.detach()) +""" + +import torch +import torch.nn as nn + + +class NegativeMemoryBank(nn.Module): + """FIFO queue of detached gallery embeddings for hard negatives. + + Args: + size: Queue capacity (number of stored embeddings). + dim: Embedding dimension. + """ + + def __init__(self, size: int = 4096, dim: int = 1024) -> None: + super().__init__() + self.size = size + self.dim = dim + # Queue stored as buffer (not a parameter, moves with .to(device)). + self.register_buffer("queue", torch.randn(size, dim)) + self.queue = nn.functional.normalize(self.queue, dim=-1) + self.register_buffer("ptr", torch.zeros(1, dtype=torch.long)) + self.register_buffer("full", torch.zeros(1, dtype=torch.bool)) + + @torch.no_grad() + def enqueue(self, embeddings: torch.Tensor) -> None: + """Add embeddings to the queue (FIFO). Oldest are overwritten.""" + batch_size = embeddings.shape[0] + ptr = int(self.ptr.item()) + + if ptr + batch_size <= self.size: + self.queue[ptr:ptr + batch_size] = embeddings.detach() + else: + # Wrap around. + overflow = (ptr + batch_size) - self.size + self.queue[ptr:] = embeddings[:batch_size - overflow].detach() + self.queue[:overflow] = embeddings[batch_size - overflow:].detach() + + new_ptr = (ptr + batch_size) % self.size + self.ptr[0] = new_ptr + if not self.full.item() and (new_ptr < ptr or new_ptr == 0): + self.full[0] = True + + def get_queue(self) -> torch.Tensor: + """Return valid queue entries [Q, dim].""" + if self.full.item(): + return self.queue + ptr = int(self.ptr.item()) + if ptr == 0: + return self.queue[:0] # empty + return self.queue[:ptr] + + @property + def current_size(self) -> int: + if self.full.item(): + return self.size + return int(self.ptr.item()) diff --git a/src/losses/multi_infonce.py b/src/losses/multi_infonce.py index 7749350..1ff4d06 100644 --- a/src/losses/multi_infonce.py +++ b/src/losses/multi_infonce.py @@ -23,14 +23,36 @@ def _symmetric_info_nce( label_smoothing: float, weight_a2b: float = 0.5, weight_b2a: float = 0.5, + queue_negatives: torch.Tensor | None = None, ) -> torch.Tensor: - """Weighted symmetric InfoNCE. Positives on the diagonal.""" + """Weighted symmetric InfoNCE with optional hard negative queue. + + Args: + emb_a: Query embeddings [B, D]. + emb_b: Gallery embeddings [B, D]. Positives on diagonal. + queue_negatives: Extra gallery negatives [Q, D] from memory bank. + """ batch_size = emb_a.size(0) - # Compute logits in fp32 to avoid overflow with small temperature. - logits = emb_a.float() @ emb_b.float().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) + emb_a_f = emb_a.float() + emb_b_f = emb_b.float() + + if queue_negatives is not None and queue_negatives.shape[0] > 0: + # a→b: query sees B in-batch + Q queue negatives. + all_b = torch.cat([emb_b_f, queue_negatives.float()], dim=0) # [B+Q, D] + logits_a2b = emb_a_f @ all_b.t() / temperature # [B, B+Q] + targets_a = torch.arange(batch_size, device=emb_a.device) + loss_a2b = F.cross_entropy(logits_a2b, targets_a, label_smoothing=label_smoothing) + + # b→a: gallery sees B in-batch queries (no queue for this direction). + logits_b2a = emb_b_f @ emb_a_f.t() / temperature # [B, B] + targets_b = torch.arange(batch_size, device=emb_a.device) + loss_b2a = F.cross_entropy(logits_b2a, targets_b, label_smoothing=label_smoothing) + else: + logits = emb_a_f @ emb_b_f.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 @@ -106,17 +128,18 @@ class InfoNCELoss(nn.Module): embeddings: dict[str, torch.Tensor], epoch: int, total_epochs: int, + queue_negatives: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: - """Compute InfoNCE loss. + """Compute InfoNCE loss with optional hard negative queue. Args: - embeddings: Dict with 'query' and 'gallery' [B, D] L2-normalized, - plus 'gate' (float) from fusion module. + embeddings: Dict with 'query' and 'gallery' [B, D] L2-normalized. epoch: Current epoch (0-indexed). total_epochs: Total epochs for temperature schedule. + queue_negatives: Extra gallery negatives [Q, D] from memory bank. Returns: - Dict with 'total', 'temperature', 'gate'. + Dict with 'total', 'temperature', 'gate_q', 'gate_g'. """ if self.learnable_temperature: # Clamp logit_scale in logit space first to prevent exp() overflow in fp16. @@ -143,6 +166,7 @@ class InfoNCELoss(nn.Module): label_smoothing=self.label_smoothing, weight_a2b=self.weight_q2g, weight_b2a=self.weight_g2q, + queue_negatives=queue_negatives, ) gate_q = embeddings.get("gate_q", embeddings.get("gate", 1.0)) diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 07b633e..9f077a6 100644 --- a/src/models/asymmetric_encoder.py +++ b/src/models/asymmetric_encoder.py @@ -297,14 +297,14 @@ class AsymmetricEncoder(nn.Module): lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt", init_gate: float = 0.7, baseline_mode: bool = False, - shared_encoder: bool = True, - embed_dim: int = 512, + shared_encoder: bool = False, mona_bottleneck: int = 64, + mona_last_n_blocks: int = 24, lora_rank: int = 4, device: str = "cuda", ) -> None: super().__init__() - self.embed_dim = embed_dim + self.embed_dim = self.DINO_DIM # native 1024, no projection self.baseline_mode = baseline_mode self.shared_encoder = shared_encoder self.device = device @@ -313,20 +313,17 @@ class AsymmetricEncoder(nn.Module): if shared_encoder: self.image_encoder = DINOv3ViT.from_pretrained(dino_web_path) self._freeze(self.image_encoder) - inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck) + inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks) LOGGER.info("Shared encoder mode: single DINOv3 WEB for drone + satellite") else: self.drone_encoder = DINOv3ViT.from_pretrained(dino_web_path) self.sat_encoder = DINOv3ViT.from_pretrained(dino_sat_path) self._freeze(self.drone_encoder) self._freeze(self.sat_encoder) - inject_mona_into_dinov3(self.drone_encoder, bottleneck=mona_bottleneck) - inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck) + inject_mona_into_dinov3(self.drone_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks) + inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks) LOGGER.info("Asymmetric encoder mode: DINOv3 WEB (drone) + DINOv3 SAT (satellite)") - # Projection: DINOv3 1024-dim -> embed_dim (512). - self.image_projection = nn.Linear(self.DINO_DIM, embed_dim) - # Text encoder — official DGTRS architecture (frozen + LoRA). if not baseline_mode: self.text_encoder = load_dgtrs_text_encoder(lrsclip_path) @@ -335,11 +332,11 @@ class AsymmetricEncoder(nn.Module): else: self.text_encoder = None - # Shared text fusion MLP: 3×768 -> embed_dim (512). + # Shared text fusion MLP: 3×768 -> 1024 (native DINOv3 dim). if not baseline_mode: self.text_fusion = TextFusionMLP( text_dim=self.TEXT_DIM, - out_dim=embed_dim, + out_dim=self.DINO_DIM, ) # Separate gated fusion for query and gallery branches. @@ -353,20 +350,16 @@ class AsymmetricEncoder(nn.Module): module.eval() def encode_drone(self, images: torch.Tensor) -> torch.Tensor: - """Encode drone images with MONA adapters + projection. Returns [B, embed_dim].""" + """Encode drone images with MONA adapters. Returns [B, 1024].""" if self.shared_encoder: - x = self.image_encoder(images) - else: - x = self.drone_encoder(images) - return self.image_projection(x) + return self.image_encoder(images) + return self.drone_encoder(images) def encode_satellite(self, images: torch.Tensor) -> torch.Tensor: - """Encode satellite images with MONA adapters + projection. Returns [B, embed_dim].""" + """Encode satellite images with MONA adapters. Returns [B, 1024].""" if self.shared_encoder: - x = self.image_encoder(images) - else: - x = self.sat_encoder(images) - return self.image_projection(x) + return self.image_encoder(images) + return self.sat_encoder(images) def encode_text_levels( self, @@ -459,7 +452,6 @@ class AsymmetricEncoder(nn.Module): "model_state": self.state_dict(), "baseline_mode": self.baseline_mode, "shared_encoder": self.shared_encoder, - "embed_dim": self.embed_dim, **extra, } tmp = path.with_suffix(path.suffix + ".tmp") @@ -494,8 +486,7 @@ class AsymmetricEncoder(nn.Module): dino_sat_path=dino_sat_path, lrsclip_path=lrsclip_path, baseline_mode=ckpt.get("baseline_mode", False), - shared_encoder=ckpt.get("shared_encoder", True), - embed_dim=ckpt.get("embed_dim", 512), + shared_encoder=ckpt.get("shared_encoder", False), device=device, ) model.load_state_dict(ckpt["model_state"], strict=False) diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 7950bf9..230bf74 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -31,6 +31,7 @@ from tqdm import tqdm from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch from src.losses.multi_infonce import InfoNCELoss +from src.losses.hard_negatives import NegativeMemoryBank from src.training.plot_metrics import generate_plots from src.training.trackers import ExperimentTracker from src.training.grad_monitor import compute_gradient_norms, log_gradient_summary @@ -73,7 +74,7 @@ class TrainConfigGTAUAV: lrsclip_path: str = _LRSCLIP init_gate: float = 0.7 baseline_mode: bool = False - shared_encoder: bool = True # single DINOv3 WEB for both branches (saves ~4-5 GB) + shared_encoder: bool = False # asymmetric: WEB (drone) + SAT (satellite) gradient_checkpointing: bool = True # trade compute for VRAM (allows larger batch) # Training. @@ -99,6 +100,7 @@ class TrainConfigGTAUAV: weight_q2g: float = 0.6 weight_g2q: float = 0.4 learnable_temperature: bool = True + neg_bank_size: int = 4096 # hard negative memory bank size (0 = disabled) # Tracking & diagnostics. use_wandb: bool = False @@ -405,6 +407,7 @@ def train(cfg: TrainConfigGTAUAV) -> None: shared_encoder=cfg.shared_encoder, device=cfg.device, ).to(cfg.device) + LOGGER.info("embed_dim=%d", model.embed_dim) # --- Gradient checkpointing (trade compute for VRAM) --- if cfg.gradient_checkpointing: @@ -446,6 +449,12 @@ def train(cfg: TrainConfigGTAUAV) -> None: cfg.tau_init, ) + # Hard negative memory bank. + neg_bank = None + if cfg.neg_bank_size > 0: + neg_bank = NegativeMemoryBank(size=cfg.neg_bank_size, dim=model.embed_dim).to(cfg.device) + LOGGER.info("Negative memory bank: size=%d, dim=%d", cfg.neg_bank_size, model.embed_dim) + # Data — separate transforms for train (augmented) and eval (clean). drone_train_tf = get_drone_train_transform(image_size=256) sat_train_tf = get_satellite_train_transform(image_size=256) @@ -599,13 +608,19 @@ def train(cfg: TrainConfigGTAUAV) -> None: sat_caption_l2=batch["sat_caption_l2"], sat_caption_l3=batch["sat_caption_l3"], ) - # Loss in fp32 (learnable temperature gradient overflows in fp16). + # Loss in fp32 with optional hard negative queue. + queue_neg = neg_bank.get_queue() if neg_bank is not None else None loss_dict = loss_fn( embeddings=embeddings, epoch=epoch, total_epochs=cfg.epochs, + queue_negatives=queue_neg, ) + # Enqueue current gallery embeddings (detached). + if neg_bank is not None: + neg_bank.enqueue(embeddings["gallery"].detach()) + # Scale loss by accumulation steps so gradients average correctly. raw_loss = float(loss_dict["total"].item()) # save before backward total_loss = loss_dict["total"] / accum