Files
caption-test/conf/gtauav_balanced.gin
pikaliov cb477f4b40 Simplify model: shared DINOv3 WEB + MONA in last 12/24 blocks
Three related architecture changes, driven by a cost/simplicity trade-off:

1. **Shared encoder**: one DINOv3 LVD-1689M (WEB) processes both drone
   and satellite images. Previously asymmetric — separate WEB (drone) and
   SAT-493M (satellite) encoders. Saves ~303M frozen params and halves
   VRAM for the image tower. Expected to lose some satellite-domain
   inductive bias; MONA adapters pick up the slack.

2. **MONA in last 12/24 blocks**: adapters injected only in the top half
   of the ViT. The lowest 12 blocks keep their pretrained features
   untouched. Trainable MONA count drops from 14.0M (48 adapters × 2
   encoders) to 3.5M (24 adapters × 1 encoder).

3. **No DINO_SAT**: `nn_models/DINO_SAT` is no longer loaded by the
   default config. It stays on disk and the path param is kept for
   backward compat with asymmetric checkpoints.

Parameter counts (with text fusion + LoRA + gates):
  Before: 17.6M trainable / 733M total (2.35%)
  After:   7.06M trainable / 434M total (1.63%)

Also fixes a pre-existing resume bug: checkpoints now record
`shared_encoder`, `baseline_mode`, `mona_bottleneck`, `mona_last_n_blocks`
so `AsymmetricEncoder.load_checkpoint` can rebuild the right architecture.
Old checkpoints still load (missing keys fall back to asymmetric defaults
via `ckpt.get(..., <default>)`).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-24 16:26:17 +03:00

65 lines
2.3 KiB
Plaintext

# GTA-UAV Balanced: Asymmetric DINOv3 (WEB+SAT) with L1/L2/L3 captions.
# Symmetric InfoNCE + MutuallyExclusiveSampler (no false negatives).
# 10 epochs, MONA all 24 blocks, 1024-dim retrieval, hard negative bank.
import src.losses.multi_infonce
# ---- Training ----
TrainConfigGTAUAV.epochs = 10
TrainConfigGTAUAV.batch_size = 8
TrainConfigGTAUAV.num_workers = 4
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 = 8
TrainConfigGTAUAV.use_amp = True
TrainConfigGTAUAV.eval_every = 1
TrainConfigGTAUAV.warmup_epochs = 2
TrainConfigGTAUAV.seed = 42
TrainConfigGTAUAV.device = "cuda"
# ---- Model ----
TrainConfigGTAUAV.init_gate = 0.7
TrainConfigGTAUAV.baseline_mode = False
TrainConfigGTAUAV.shared_encoder = True # single DINOv3 WEB for both branches
TrainConfigGTAUAV.mona_bottleneck = 64
TrainConfigGTAUAV.mona_last_n_blocks = 12 # inject MONA only in last 12/24 ViT blocks
TrainConfigGTAUAV.gradient_checkpointing = True
# ---- Loss ----
TrainConfigGTAUAV.loss_type = "symmetric"
TrainConfigGTAUAV.tau_init = 0.07
TrainConfigGTAUAV.label_smoothing = 0.1
TrainConfigGTAUAV.learnable_temperature = True
TrainConfigGTAUAV.weight_q2g = 0.6
TrainConfigGTAUAV.weight_g2q = 0.4
TrainConfigGTAUAV.neg_bank_size = 4096
# ---- Sampling ----
TrainConfigGTAUAV.sampler_type = "dss" # "dss" or "mutex"
TrainConfigGTAUAV.dss_warmup_epochs = 1 # first N epochs use mutex-only (untrained embeds not useful)
TrainConfigGTAUAV.dss_reembed_every = 1
TrainConfigGTAUAV.use_mutex_sampler = True # legacy flag, kept True unless disabling both samplers
# ---- Output ----
TrainConfigGTAUAV.output_dir = "out/gtauav/with_text"
# ---- Tracking ----
TrainConfigGTAUAV.use_wandb = False
TrainConfigGTAUAV.use_tb = True
TrainConfigGTAUAV.use_gradcam = True
TrainConfigGTAUAV.gradcam_every = 5
TrainConfigGTAUAV.use_profiler = False
TrainConfigGTAUAV.log_grad_norms = True
# ---- InfoNCELoss (gin-configurable) ----
InfoNCELoss.temperature_init = 0.07
InfoNCELoss.learnable_temperature = True
InfoNCELoss.label_smoothing = 0.1
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)