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

Asymmetric encoders + MONA all 24 blocks + 1024-dim + hard negatives

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-24 08:47:33 +03:00
parent 75a4350d18
commit 04d5307221
5 changed files with 142 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
conf/gtauav_balanced.gin
View File

@@ -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"

70
src/losses/hard_negatives.py Normal file
View File

@@ -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())

44
src/losses/multi_infonce.py
View File

@@ -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))

39
src/models/asymmetric_encoder.py
View File

@@ -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)

19
src/training/train_gtauav.py
View File

@@ -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