diff --git a/CLAUDE.md b/CLAUDE.md index 5028961..18b0884 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -101,7 +101,7 @@ Eval: Resize(256) + CenterCrop(256) + ImageNet normalization. | `src/models/asymmetric_encoder.py` | DINOv3ViT + TextFusionMLP + AsymmetricEncoder + GatedFusion + encode_query/encode_gallery (per-sample caption mask) | | `src/datasets/gtauav_dataset.py` | GTA-UAV-LR loader + L1/L2/L3 captions + GTAUAVSatGallery/GTAUAVDroneQuery (full retrieval eval) | | `src/datasets/mutually_exclusive_sampler.py` | BatchSampler: drone'ы в батче не делят sat_candidates (no false negatives) | -| `src/losses/multi_infonce.py` | **Primary:** SymmetricInfoNCE + MoCo queue, learnable τ clamp [0.01, 0.1], weights q2g=0.6 g2q=0.4 | +| `src/losses/multi_infonce.py` | **Primary:** SymmetricInfoNCE + MoCo queue, learnable τ clamp [0.01, 0.1], weights q2g=0.6 g2q=0.4, `hard_mining_k` для top-K hardest negatives | | `src/losses/weighted_infonce.py` | Alternative: per-sample adaptive label smoothing (активируется `loss_type="weighted"`) | | `src/losses/hard_negatives.py` | NegativeMemoryBank (MoCo-style FIFO queue 4096 × 1024) | | `src/training/train_gtauav.py` | Training loop: full-gallery `_evaluate`, mutex sampler wiring, loss_type switch | @@ -212,6 +212,7 @@ Meta-файл `meta/seg_filter.json`: исключение изображени - **Optimizer:** AdamW, per-group LR: proj=1e-4, text=1e-5 (10x lower) - **Scheduler:** linear warmup (2 epochs) + cosine annealing (per-step) - **Loss:** SymmetricInfoNCE (q2g=0.6, g2q=0.4) с learnable τ (init=0.07, clamp [0.01, 0.1]) +- **Hard mining:** top-K=512 hardest negatives per query из MoCo queue (размер 4096); `hard_mining_k=0` отключает - **Batch sampler:** MutuallyExclusiveSampler — batches disjoint по sat_candidates (на bs=8 сохраняет 100% entries) - **Eval:** full satellite gallery (~2684 unique tiles для test_20) с multi-match R@K (учитывает все positive/semi-positive) - **Augmentations:** diff --git a/conf/gtauav_balanced.gin b/conf/gtauav_balanced.gin index 9461348..0d756e2 100644 --- a/conf/gtauav_balanced.gin +++ b/conf/gtauav_balanced.gin @@ -56,3 +56,4 @@ InfoNCELoss.weight_q2g = 0.6 InfoNCELoss.weight_g2q = 0.4 InfoNCELoss.tau_min = 0.01 InfoNCELoss.tau_max = 0.1 +InfoNCELoss.hard_mining_k = 512 # 0 = use whole queue (disable mining) diff --git a/scripts/smoke_train.py b/scripts/smoke_train.py index 3270943..3be0c14 100644 --- a/scripts/smoke_train.py +++ b/scripts/smoke_train.py @@ -8,6 +8,7 @@ from torch.utils.data import DataLoader from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler +from src.losses.hard_negatives import NegativeMemoryBank from src.losses.multi_infonce import InfoNCELoss from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform @@ -40,13 +41,15 @@ def main() -> None: temperature_init=0.07, learnable_temperature=True, label_smoothing=0.1, weight_q2g=0.6, weight_g2q=0.4, tau_min=0.01, tau_max=0.1, + hard_mining_k=512, ).to("cuda") + neg_bank = NegativeMemoryBank(size=4096, dim=model.embed_dim).to("cuda") trainable = [p for p in model.trainable_parameters()] + list(loss_fn.parameters()) opt = torch.optim.AdamW(trainable, lr=1e-4) it = iter(loader) - for step in range(2): + for step in range(3): batch = next(it) opt.zero_grad() emb = model( @@ -55,9 +58,11 @@ def main() -> None: caption_l1=batch["caption_l1"], caption_l2=batch["caption_l2"], caption_l3=batch["caption_l3"], sat_caption_l1=batch["sat_caption_l1"], sat_caption_l2=batch["sat_caption_l2"], sat_caption_l3=batch["sat_caption_l3"], ) - out = loss_fn(emb, epoch=0, total_epochs=10) + queue = neg_bank.get_queue() + out = loss_fn(emb, epoch=0, total_epochs=10, queue_negatives=queue) out["total"].backward() opt.step() + neg_bank.enqueue(emb["gallery"].detach()) # Verify mutual exclusion in batch batch_sats = [set(ds.entries[i]["sat_candidates"]) for i in batch.get("__indices__", range(8))] @@ -67,12 +72,12 @@ def main() -> None: f" step {step}: loss={out['total'].item():.4f} " f"tau={out['temperature'].item():.4f} " f"gate_q={out['gate_q'].item():.3f} gate_g={out['gate_g'].item():.3f} " - f"n_drone_caps={sum(1 for t in batch['caption_l1'] if t)} " - f"n_sat_caps={sum(1 for t in batch['sat_caption_l1'] if t)}" + f"queue_size={queue.shape[0] if queue is not None else 0} " + f"mining_k={loss_fn.hard_mining_k if queue is not None and queue.shape[0] > loss_fn.hard_mining_k else 'full'}" ) assert torch.isfinite(out["total"]).all(), "Loss not finite!" - print("OK: 2 train steps completed with finite loss") + print("OK: 3 train steps completed with finite loss (hard mining K=512)") if __name__ == "__main__": diff --git a/src/losses/multi_infonce.py b/src/losses/multi_infonce.py index 87a3449..445e9a0 100644 --- a/src/losses/multi_infonce.py +++ b/src/losses/multi_infonce.py @@ -24,6 +24,7 @@ def _symmetric_info_nce( weight_a2b: float = 0.5, weight_b2a: float = 0.5, queue_negatives: torch.Tensor | None = None, + hard_mining_k: int = 0, ) -> torch.Tensor: """Weighted symmetric InfoNCE with optional hard negative queue. @@ -31,20 +32,31 @@ def _symmetric_info_nce( 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. + hard_mining_k: If > 0 and queue is non-empty, use only the top-K + hardest (highest-similarity) queue entries per query instead + of the full queue. Per-query selection — each row gets its + own K negatives gathered via `topk`. """ batch_size = emb_a.size(0) 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] + queue_f = queue_negatives.float() + sim_inbatch = emb_a_f @ emb_b_f.t() / temperature # [B, B] + sim_queue = emb_a_f @ queue_f.t() / temperature # [B, Q] + + if hard_mining_k > 0 and hard_mining_k < queue_f.shape[0]: + # Per-row top-K — each query gets its own hardest negatives. + sim_queue, _ = sim_queue.topk(k=hard_mining_k, dim=1) # [B, K] + + # a→b: [B, B + (Q or K)]. Positive at column `i` for row `i`. + logits_a2b = torch.cat([sim_inbatch, sim_queue], dim=1) 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] + # b→a: gallery sees B in-batch queries (queue is gallery-side, irrelevant here). + logits_b2a = sim_inbatch.t() # [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: @@ -83,6 +95,9 @@ class InfoNCELoss(nn.Module): (CLIP-style logit_scale). If False, uses cosine schedule. tau_min: Minimum clamp for learnable temperature. tau_max: Maximum clamp for learnable temperature. + hard_mining_k: If > 0, mine top-K hardest negatives per query from + the memory bank queue instead of using the full queue. 0 disables + mining (queue used whole). Typical values: 256-1024 for queue=4096. """ def __init__( @@ -95,6 +110,7 @@ class InfoNCELoss(nn.Module): learnable_temperature: bool = True, tau_min: float = 0.01, tau_max: float = 0.1, + hard_mining_k: int = 0, ) -> None: super().__init__() self.temperature_init = temperature_init @@ -105,6 +121,7 @@ class InfoNCELoss(nn.Module): self.learnable_temperature = learnable_temperature self.tau_min = tau_min self.tau_max = tau_max + self.hard_mining_k = hard_mining_k if learnable_temperature: # Store as log(1/tau) like CLIP's logit_scale. @@ -167,6 +184,7 @@ class InfoNCELoss(nn.Module): weight_a2b=self.weight_q2g, weight_b2a=self.weight_g2q, queue_negatives=queue_negatives, + hard_mining_k=self.hard_mining_k, ) gate_q = embeddings.get("gate_q", embeddings.get("gate", 1.0))