diff --git a/CLAUDE.md b/CLAUDE.md index 47d29b4..49875a3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -97,7 +97,16 @@ Eval: Resize(256) + CenterCrop(256) + ImageNet normalization. | `src/models/asymmetric_encoder.py` | DINOv3ViT + TextFusionMLP + AsymmetricEncoder + GatedFusion | | `src/datasets/gtauav_dataset.py` | GTA-UAV-LR loader + L1/L2/L3 caption parsing из VLM JSON | | `src/losses/multi_infonce.py` | InfoNCE с learnable temperature (fp32), clamp [0.01, 0.5] | -| `src/training/train_gtauav.py` | Training loop с eval, AMP, per-group LR, warmup, --resume | +| `src/training/train_gtauav.py` | Training loop с gin, W&B/TB, AMP, per-group LR, warmup, --resume | +| `src/training/trackers.py` | Unified experiment tracker: W&B + TensorBoard + CSV | +| `src/training/grad_monitor.py` | Gradient norm monitoring per param group | +| `src/training/gradcam.py` | Grad-CAM visualization для DINOv3 encoders | +| `src/training/profiling.py` | PyTorch Profiler wrapper + torchinfo model summary | +| `src/training/plot_metrics.py` | Seaborn/matplotlib plots (каждую эпоху) | +| `conf/gtauav_balanced.gin` | With text, gate=0.7, 10 epochs | +| `conf/gtauav_baseline.gin` | No text, gate=1.0 | +| `conf/gtauav_text_heavy.gin` | Text-heavy, gate=0.3 | +| `conf/gtauav_image_heavy.gin` | Image-heavy, gate=0.9 | | `scripts/make_split.py` | 80/20 random split из всех пар | | `scripts/filter_segmentation.py` | Scan segm masks, output meta JSON (exclude >=90% bg+water) | @@ -204,6 +213,14 @@ Meta-файл `meta/seg_filter.json`: исключение изображени - Test: 6,742 → 6,252 after seg filter - Скрипт: `python -m scripts.make_split --ratio 0.8 --seed 42` +### V3 (GTA-UAV, gin) +| Конфиг | Gate init | Описание | +|--------|-----------|----------| +| `conf/gtauav_balanced.gin` | 0.7 (30% text) | **Primary test** | +| `conf/gtauav_baseline.gin` | 1.0 (no text) | Reference baseline | +| `conf/gtauav_text_heavy.gin` | 0.3 (70% text) | Stress test | +| `conf/gtauav_image_heavy.gin` | 0.9 (10% text) | Image-dominant | + ### V2 (UAV-GeoLoc, gin) | Конфиг | Gate init | Описание | |--------|-----------|----------| @@ -218,17 +235,31 @@ Meta-файл `meta/seg_filter.json`: исключение изображени # 1. Filter segmentation (exclude 90%+ background/water) python -m scripts.filter_segmentation --output meta/seg_filter.json -# 2. Baseline (no text) -python -m src.training.train_gtauav --baseline --filter-meta meta/seg_filter.json +# 2. Train with gin config (recommended) +python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \ + --filter-meta meta/seg_filter.json -# 3. With captions (L1/L2/L3) -python -m src.training.train_gtauav --filter-meta meta/seg_filter.json +# 3. Baseline (no text) +python -m src.training.train_gtauav --config conf/gtauav_baseline.gin \ + --filter-meta meta/seg_filter.json -# 4. Compare +# 4. With diagnostics (W&B + Grad-CAM + Profiler) +python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \ + --filter-meta meta/seg_filter.json --wandb --gradcam --profile + +# 5. CLI overrides (gin params take priority) +python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \ + --filter-meta meta/seg_filter.json \ + --gin-param 'TrainConfigGTAUAV.batch_size=16' + +# 6. Compare python -m scripts.compare_runs \ --baseline_report out/gtauav/baseline/eval_report.json \ --full_report out/gtauav/with_text/eval_report.json \ --output out/gtauav/comparison.md + +# 7. TensorBoard +tensorboard --logdir out/gtauav/with_text/tb_logs ``` ### V2 (UAV-GeoLoc) diff --git a/README.md b/README.md index a6e0965..67a3116 100644 --- a/README.md +++ b/README.md @@ -218,10 +218,13 @@ Gallery: sat_img -> GeoRSCLIP -> gallery ``` caption-test/ -├── conf/ # Gin configs (v2) -│ ├── balanced.gin -│ ├── baseline_no_text.gin -│ └── text_heavy.gin +├── conf/ # Gin configs +│ ├── gtauav_balanced.gin # GTA-UAV with text (10 epochs, v3) +│ ├── gtauav_baseline.gin # GTA-UAV baseline, no text (v3) +│ ├── gtauav_text_heavy.gin # GTA-UAV text-heavy gate=0.3 (v3) +│ ├── balanced.gin # UAV-GeoLoc with text (v2) +│ ├── baseline_no_text.gin # UAV-GeoLoc baseline (v2) +│ └── text_heavy.gin # UAV-GeoLoc text-heavy (v2) ├── nn_models/ # Pre-trained checkpoints (v3, gitignored) │ ├── DINO_WEB/ # DINOv3 ViT-L/16 LVD-1689M (.pth) │ ├── DINO_SAT/ # DINOv3 ViT-L/16 SAT-493M (.safetensors) @@ -250,7 +253,12 @@ caption-test/ │ │ └── multi_infonce.py # InfoNCE with learnable temperature │ ├── training/ │ │ ├── train_gtauav.py # Training loop GTA-UAV (v3) -│ │ └── train.py # Training loop UAV-GeoLoc (v2) +│ │ ├── train.py # Training loop UAV-GeoLoc (v2) +│ │ ├── trackers.py # Unified tracker: W&B + TensorBoard +│ │ ├── grad_monitor.py # Gradient norm monitoring per group +│ │ ├── gradcam.py # Grad-CAM visualization for DINOv3 +│ │ ├── profiling.py # PyTorch Profiler + torchinfo summary +│ │ └── plot_metrics.py # Seaborn/matplotlib metric plots │ └── eval/ │ └── evaluate.py # R@K metrics, Delta R@1 └── checkpoints/ # GeoRSCLIP RS5M_ViT-B-32.pt (v2) @@ -268,6 +276,17 @@ regex gin-config Pillow numpy +pandas +matplotlib +seaborn +``` + +### Optional (for extended diagnostics) + +``` +wandb # Weights & Biases experiment tracking +torchinfo # Model summary tables +tensorboard # TensorBoard logging (included with torch) ``` ## Workflow (V3 — GTA-UAV) @@ -279,26 +298,52 @@ python -m scripts.make_split --output-dir meta python -m scripts.filter_segmentation --output meta/seg_filter.json ``` -### 2. Train baseline (no text) +### 2. Train with gin configs (recommended) + +```bash +# Baseline (no text, 10 epochs) +python -m src.training.train_gtauav --config conf/gtauav_baseline.gin \ + --filter-meta meta/seg_filter.json + +# With captions (L1/L2/L3, 10 epochs) +python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \ + --filter-meta meta/seg_filter.json + +# Text-heavy (gate=0.3, 70% text weight) +python -m src.training.train_gtauav --config conf/gtauav_text_heavy.gin \ + --filter-meta meta/seg_filter.json +``` + +### 3. Train without gin (CLI-only) ```bash python -m src.training.train_gtauav --baseline --filter-meta meta/seg_filter.json -``` - -### 3. Train with captions (L1/L2/L3) - -```bash python -m src.training.train_gtauav --filter-meta meta/seg_filter.json ``` -### 4. Resume from checkpoint +### 4. Enable diagnostics + +```bash +# W&B + Grad-CAM + PyTorch Profiler +python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \ + --filter-meta meta/seg_filter.json --wandb --gradcam --profile + +# Gin parameter overrides from CLI +python -m src.training.train_gtauav --config conf/gtauav_balanced.gin \ + --filter-meta meta/seg_filter.json \ + --gin-param 'TrainConfigGTAUAV.batch_size=16' 'TrainConfigGTAUAV.epochs=20' +``` + +CLI flags (`--wandb`, `--gradcam`, `--profile`, `--epochs`, etc.) take priority over gin config. + +### 5. Resume from checkpoint ```bash python -m src.training.train_gtauav --resume out/gtauav/with_text/ckpt_epoch004.pt \ --filter-meta meta/seg_filter.json ``` -### 5. Compare and get verdict +### 6. Compare and get verdict ```bash python -m scripts.compare_runs \ @@ -307,6 +352,24 @@ python -m scripts.compare_runs \ --output out/gtauav/comparison.md ``` +### 7. View TensorBoard + +```bash +tensorboard --logdir out/gtauav/with_text/tb_logs +``` + +## Diagnostics & Visualization + +| Tool | Flag | Output | Description | +|------|------|--------|-------------| +| **TensorBoard** | `--use-tb` (default on) | `{out}/tb_logs/` | Scalars, histograms, images | +| **W&B** | `--wandb` | cloud | Full experiment tracking, Grad-CAM images | +| **Grad-CAM** | `--gradcam` | `{out}/gradcam/` | DINOv3 attention heatmaps (drone + satellite) | +| **PyTorch Profiler** | `--profile` | `{out}/profiler/` | Chrome trace, CUDA timeline, memory | +| **torchinfo** | auto | `{out}/model_summary.txt` | Layer-by-layer parameter table | +| **Gradient norms** | `--log-grad-norms` (default on) | TB/W&B | Per-group: MONA, LoRA, MLP, gates, tau | +| **CSV + plots** | auto | `{out}/logs/` | train.csv, val.csv, PNG plots every epoch | + ## Decision rule | Delta R@1 (drone→satellite) | Verdict | diff --git a/conf/gtauav_balanced.gin b/conf/gtauav_balanced.gin new file mode 100644 index 0000000..ce3075e --- /dev/null +++ b/conf/gtauav_balanced.gin @@ -0,0 +1,50 @@ +# GTA-UAV Balanced: GatedFusion with L1/L2/L3 captions on both branches. +# query = sigma(alpha) * drone + (1-sigma(alpha)) * text -> InfoNCE vs gallery +# 20 epochs, DINOv3 + DGTRS-CLIP, MONA + LoRA adapters. + +import src.losses.multi_infonce +import src.training.train_gtauav + +# ---- 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.use_amp = True +TrainConfigGTAUAV.eval_every = 2 +TrainConfigGTAUAV.warmup_epochs = 2 +TrainConfigGTAUAV.seed = 42 +TrainConfigGTAUAV.device = "cuda" + +# ---- Model ---- +TrainConfigGTAUAV.init_gate = 0.7 +TrainConfigGTAUAV.baseline_mode = False + +# ---- Loss ---- +TrainConfigGTAUAV.tau_init = 0.07 +TrainConfigGTAUAV.label_smoothing = 0.1 +TrainConfigGTAUAV.weight_q2g = 0.6 +TrainConfigGTAUAV.weight_g2q = 0.4 +TrainConfigGTAUAV.learnable_temperature = True + +# ---- 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 + +# ---- InfoNCE Loss (gin-configurable) ---- +InfoNCELoss.temperature_init = 0.07 +InfoNCELoss.temperature_final = 0.01 +InfoNCELoss.label_smoothing = 0.1 +InfoNCELoss.weight_q2g = 0.6 +InfoNCELoss.weight_g2q = 0.4 +InfoNCELoss.learnable_temperature = True diff --git a/conf/gtauav_baseline.gin b/conf/gtauav_baseline.gin new file mode 100644 index 0000000..3bcadde --- /dev/null +++ b/conf/gtauav_baseline.gin @@ -0,0 +1,9 @@ +# GTA-UAV Baseline: no text fusion (gate forced to 1.0). +# query = drone_only -> InfoNCE vs satellite +# Reference R@1 for delta computation. + +include 'conf/gtauav_balanced.gin' + +TrainConfigGTAUAV.baseline_mode = True +TrainConfigGTAUAV.output_dir = "out/gtauav/baseline" +TrainConfigGTAUAV.use_gradcam = False diff --git a/conf/gtauav_image_heavy.gin b/conf/gtauav_image_heavy.gin new file mode 100644 index 0000000..93bb699 --- /dev/null +++ b/conf/gtauav_image_heavy.gin @@ -0,0 +1,8 @@ +# GTA-UAV Image-heavy: gate initialized high (more image weight). +# query = sigma(0.9) * drone + 0.1 * text +# Minimal text contribution test. + +include 'conf/gtauav_balanced.gin' + +TrainConfigGTAUAV.init_gate = 0.9 +TrainConfigGTAUAV.output_dir = "out/gtauav/image_heavy" diff --git a/conf/gtauav_text_heavy.gin b/conf/gtauav_text_heavy.gin new file mode 100644 index 0000000..3791b59 --- /dev/null +++ b/conf/gtauav_text_heavy.gin @@ -0,0 +1,8 @@ +# GTA-UAV Text-heavy: gate initialized low (more text weight). +# query = sigma(0.3) * drone + 0.7 * text +# Stress test for text contribution. + +include 'conf/gtauav_balanced.gin' + +TrainConfigGTAUAV.init_gate = 0.3 +TrainConfigGTAUAV.output_dir = "out/gtauav/text_heavy" diff --git a/src/training/grad_monitor.py b/src/training/grad_monitor.py new file mode 100644 index 0000000..dfcaebe --- /dev/null +++ b/src/training/grad_monitor.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +"""Gradient monitoring for training diagnostics. + +Computes per-group gradient norms after backward pass to detect +vanishing/exploding gradients in different components: + - MONA adapters (drone + satellite) + - LoRA (DGTRS-CLIP text encoder) + - TextFusionMLP + - GatedFusion gates (alpha_q, alpha_g) + - Learnable temperature (logit_scale) +""" + +import logging + +import torch +import torch.nn as nn + +LOGGER = logging.getLogger("caption_test.grad_monitor") + +# Parameter group patterns for classification. +_GROUP_PATTERNS = { + "mona_drone": "drone_encoder", + "mona_sat": "sat_encoder", + "lora_text": "text_encoder", + "text_fusion_mlp": "text_fusion", + "gate_q": "fusion_query", + "gate_g": "fusion_gallery", +} + + +def compute_gradient_norms( + model: nn.Module, + loss_fn: nn.Module | None = None, +) -> dict[str, float]: + """Compute L2 gradient norms per parameter group. + + Args: + model: The model after backward(). + loss_fn: Loss module (for logit_scale gradient). + + Returns: + Dict mapping group name to gradient L2 norm. + """ + group_grads: dict[str, list[torch.Tensor]] = {k: [] for k in _GROUP_PATTERNS} + group_grads["other"] = [] + + for name, param in model.named_parameters(): + if not param.requires_grad or param.grad is None: + continue + + matched = False + for group_name, pattern in _GROUP_PATTERNS.items(): + if pattern in name: + group_grads[group_name].append(param.grad.detach().flatten()) + matched = True + break + if not matched: + group_grads["other"].append(param.grad.detach().flatten()) + + norms: dict[str, float] = {} + for group_name, grads in group_grads.items(): + if grads: + all_grads = torch.cat(grads) + norms[f"grad_norm/{group_name}"] = float(all_grads.norm(2).item()) + norms[f"grad_max/{group_name}"] = float(all_grads.abs().max().item()) + + # Logit scale gradient (from loss module). + if loss_fn is not None and hasattr(loss_fn, "logit_scale"): + ls = loss_fn.logit_scale + if ls is not None and ls.grad is not None: + norms["grad_norm/logit_scale"] = float(ls.grad.detach().abs().item()) + + return norms + + +def log_gradient_summary(norms: dict[str, float]) -> None: + """Log gradient norms summary to console.""" + norm_parts = [] + for k, v in sorted(norms.items()): + if k.startswith("grad_norm/"): + name = k.replace("grad_norm/", "") + norm_parts.append(f"{name}={v:.4f}") + if norm_parts: + LOGGER.info("grad norms: %s", " ".join(norm_parts)) diff --git a/src/training/gradcam.py b/src/training/gradcam.py new file mode 100644 index 0000000..c9d2e91 --- /dev/null +++ b/src/training/gradcam.py @@ -0,0 +1,241 @@ +from __future__ import annotations + +"""Grad-CAM visualization for DINOv3 ViT encoders. + +Generates attention heatmaps overlaid on input images to show +which regions the model focuses on for drone/satellite matching. + +Hook target: last DINOv3Block output (before final LayerNorm). +Uses patch tokens (excluding CLS + registers) to build spatial map. + +Usage: + cam = DINOv3GradCAM(model.drone_encoder, target_layer_idx=-1) + heatmap = cam.generate(image_tensor, class_idx=None) # [H, W] numpy + overlay = cam.overlay(image_tensor, heatmap) # [3, H, W] torch + cam.remove_hooks() +""" + +import logging +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + +LOGGER = logging.getLogger("caption_test.gradcam") + + +class DINOv3GradCAM: + """Grad-CAM for DINOv3 ViT encoder. + + Hooks into a specific transformer block to capture activations + and gradients, then produces spatial attention heatmaps. + + Args: + encoder: DINOv3ViT model instance. + target_layer_idx: Which block to hook (-1 = last block). + num_registers: Number of register tokens (skipped in spatial map). + """ + + def __init__( + self, + encoder: nn.Module, + target_layer_idx: int = -1, + num_registers: int = 4, + ) -> None: + self.encoder = encoder + self.num_registers = num_registers + self._activations: torch.Tensor | None = None + self._gradients: torch.Tensor | None = None + + # Hook into the target block. + target_block = encoder.layer[target_layer_idx] + self._fwd_hook = target_block.register_forward_hook(self._save_activation) + self._bwd_hook = target_block.register_full_backward_hook(self._save_gradient) + + def _save_activation(self, module: nn.Module, input: Any, output: torch.Tensor) -> None: + self._activations = output.detach() + + def _save_gradient(self, module: nn.Module, grad_input: Any, grad_output: tuple) -> None: + self._gradients = grad_output[0].detach() + + @torch.enable_grad() + def generate( + self, + image: torch.Tensor, + target_embedding: torch.Tensor | None = None, + ) -> np.ndarray: + """Generate Grad-CAM heatmap for a single image. + + Args: + image: Input image [1, 3, H, W]. + target_embedding: Optional target to maximize similarity against. + If None, uses the CLS token norm as the scalar output. + + Returns: + Heatmap [H_patches, W_patches] as numpy array, values in [0, 1]. + """ + self.encoder.zero_grad() + image = image.requires_grad_(True) + + # Forward through the full encoder. + cls_token = self.encoder(image) # [1, D] + + # Scalar to backprop from. + if target_embedding is not None: + score = (cls_token * target_embedding).sum() + else: + score = cls_token.norm(dim=-1).sum() + + score.backward(retain_graph=True) + + if self._activations is None or self._gradients is None: + LOGGER.warning("No activations/gradients captured") + return np.zeros((16, 16), dtype=np.float32) + + # Gradients: [1, N_tokens, D] — average over token dim for weights. + weights = self._gradients.mean(dim=1, keepdim=True) # [1, 1, D] + + # Weighted activation map. + # Skip CLS (idx 0) + registers (idx 1..num_registers). + n_special = 1 + self.num_registers + patch_activations = self._activations[:, n_special:, :] # [1, N_patches, D] + + cam = (patch_activations * weights).sum(dim=-1) # [1, N_patches] + cam = F.relu(cam) # Only positive contributions. + + # Reshape to spatial grid. + n_patches = cam.shape[1] + h = w = int(n_patches ** 0.5) + cam = cam.reshape(1, 1, h, w) + + # Normalize to [0, 1]. + cam = cam - cam.min() + cam_max = cam.max() + if cam_max > 0: + cam = cam / cam_max + + return cam.squeeze().cpu().numpy() + + def overlay( + self, + image: torch.Tensor, + heatmap: np.ndarray, + alpha: float = 0.5, + colormap: str = "jet", + ) -> torch.Tensor: + """Overlay heatmap on image. + + Args: + image: Original image [1, 3, H, W] or [3, H, W] (ImageNet normalized). + heatmap: Grad-CAM heatmap [h, w] in [0, 1]. + alpha: Overlay transparency. + colormap: Matplotlib colormap name. + + Returns: + Overlaid image [3, H, W] as torch tensor, values in [0, 1]. + """ + import matplotlib.cm as cm + + if image.dim() == 4: + image = image.squeeze(0) + + _, H, W = image.shape + + # Denormalize from ImageNet stats. + mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1).to(image.device) + std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1).to(image.device) + img = (image * std + mean).clamp(0, 1).cpu().numpy().transpose(1, 2, 0) # [H, W, 3] + + # Resize heatmap to image size. + heatmap_resized = torch.tensor(heatmap).unsqueeze(0).unsqueeze(0).float() + heatmap_resized = F.interpolate(heatmap_resized, size=(H, W), mode="bilinear", align_corners=False) + heatmap_np = heatmap_resized.squeeze().numpy() + + # Apply colormap. + cmap = cm.get_cmap(colormap) + heatmap_colored = cmap(heatmap_np)[:, :, :3] # [H, W, 3] + + # Blend. + overlay = alpha * heatmap_colored + (1 - alpha) * img + overlay = np.clip(overlay, 0, 1) + + return torch.from_numpy(overlay.transpose(2, 0, 1)).float() + + def remove_hooks(self) -> None: + """Remove forward/backward hooks.""" + self._fwd_hook.remove() + self._bwd_hook.remove() + + +def generate_gradcam_samples( + model: nn.Module, + dataloader: torch.utils.data.DataLoader, + device: str, + output_dir: str, + n_samples: int = 8, + epoch: int = 0, +) -> list[torch.Tensor]: + """Generate Grad-CAM visualizations for a batch of samples. + + Creates overlaid heatmaps for both drone and satellite encoders. + Saves to output_dir/gradcam/epoch_{N}/. + + Args: + model: AsymmetricEncoder instance. + dataloader: Validation dataloader. + device: Torch device. + output_dir: Base output directory. + n_samples: Number of samples to visualize. + epoch: Current epoch (for naming). + + Returns: + List of overlay tensors [3, H, W] for logging. + """ + from pathlib import Path + import matplotlib.pyplot as plt + + save_dir = Path(output_dir) / "gradcam" / f"epoch_{epoch:03d}" + save_dir.mkdir(parents=True, exist_ok=True) + + model.eval() + cam_drone = DINOv3GradCAM(model.drone_encoder, target_layer_idx=-1) + cam_sat = DINOv3GradCAM(model.sat_encoder, target_layer_idx=-1) + + overlays = [] + batch = next(iter(dataloader)) + drone_imgs = batch["drone_img"][:n_samples].to(device) + sat_imgs = batch["sat_img"][:n_samples].to(device) + + for i in range(min(n_samples, drone_imgs.shape[0])): + drone_img = drone_imgs[i:i+1] + sat_img = sat_imgs[i:i+1] + + # Drone Grad-CAM. + heatmap_d = cam_drone.generate(drone_img) + overlay_d = cam_drone.overlay(drone_img, heatmap_d) + overlays.append(overlay_d) + + # Satellite Grad-CAM. + heatmap_s = cam_sat.generate(sat_img) + overlay_s = cam_sat.overlay(sat_img, heatmap_s) + overlays.append(overlay_s) + + # Save as side-by-side figure. + fig, axes = plt.subplots(1, 2, figsize=(10, 5)) + axes[0].imshow(overlay_d.permute(1, 2, 0).numpy()) + axes[0].set_title(f"Drone #{i}") + axes[0].axis("off") + axes[1].imshow(overlay_s.permute(1, 2, 0).numpy()) + axes[1].set_title(f"Satellite #{i}") + axes[1].axis("off") + plt.tight_layout() + fig.savefig(save_dir / f"sample_{i:03d}.png", dpi=150, bbox_inches="tight") + plt.close(fig) + + cam_drone.remove_hooks() + cam_sat.remove_hooks() + + LOGGER.info("Grad-CAM: saved %d samples to %s", len(overlays) // 2, save_dir) + return overlays diff --git a/src/training/profiling.py b/src/training/profiling.py new file mode 100644 index 0000000..989ee6f --- /dev/null +++ b/src/training/profiling.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +"""PyTorch Profiler wrapper for training performance analysis. + +Profiles the first N batches of training to identify bottlenecks +in CUDA/CPU execution, memory allocation, and data loading. + +Exports: + - Chrome trace (viewable in chrome://tracing) + - TensorBoard plugin data (if TB available) + - Summary table to console + +Usage: + profiler = TrainingProfiler(output_dir, n_warmup=3, n_active=5) + for batch_idx, batch in enumerate(loader): + with profiler.step_context(batch_idx): + # ... training step ... + if profiler.is_done(batch_idx): + break + profiler.export() +""" + +import logging +from pathlib import Path + +import torch +from torch.profiler import ProfilerActivity, profile, schedule, tensorboard_trace_handler + +LOGGER = logging.getLogger("caption_test.profiler") + + +class TrainingProfiler: + """PyTorch profiler for first N training batches. + + Args: + output_dir: Directory for profiler output. + n_warmup: Number of warmup steps (not profiled). + n_active: Number of steps to actively profile. + n_repeat: Number of profiling cycles. + record_shapes: Record tensor shapes. + profile_memory: Track memory allocation. + with_stack: Record Python call stacks. + """ + + def __init__( + self, + output_dir: str | Path, + n_warmup: int = 3, + n_active: int = 5, + n_repeat: int = 1, + record_shapes: bool = True, + profile_memory: bool = True, + with_stack: bool = False, + ) -> None: + self.output_dir = Path(output_dir) / "profiler" + self.output_dir.mkdir(parents=True, exist_ok=True) + self.n_warmup = n_warmup + self.n_active = n_active + self.n_repeat = n_repeat + self.total_steps = (n_warmup + n_active) * n_repeat + + self._profiler = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + schedule=schedule( + wait=0, + warmup=n_warmup, + active=n_active, + repeat=n_repeat, + ), + on_trace_ready=tensorboard_trace_handler(str(self.output_dir)), + record_shapes=record_shapes, + profile_memory=profile_memory, + with_stack=with_stack, + ) + self._started = False + + def start(self) -> None: + """Start the profiler.""" + self._profiler.__enter__() + self._started = True + LOGGER.info( + "Profiler started: %d warmup + %d active steps, output: %s", + self.n_warmup, self.n_active, self.output_dir, + ) + + def step(self) -> None: + """Signal end of a profiling step.""" + if self._started: + self._profiler.step() + + def is_done(self, batch_idx: int) -> bool: + """Check if profiling is complete.""" + return batch_idx >= self.total_steps + + def export(self) -> None: + """Export profiling results and print summary.""" + if not self._started: + return + + self._profiler.__exit__(None, None, None) + self._started = False + + # Print key averages summary. + summary = self._profiler.key_averages().table( + sort_by="cuda_time_total", row_limit=20, + ) + LOGGER.info("Profiler summary (top 20 by CUDA time):\n%s", summary) + + # Export Chrome trace. + trace_path = self.output_dir / "chrome_trace.json" + self._profiler.export_chrome_trace(str(trace_path)) + LOGGER.info("Chrome trace exported: %s", trace_path) + + # Memory summary if available. + if torch.cuda.is_available(): + mem_summary = torch.cuda.memory_summary(abbreviated=True) + summary_path = self.output_dir / "memory_summary.txt" + summary_path.write_text(mem_summary) + LOGGER.info("CUDA memory summary: %s", summary_path) + + +def print_model_summary(model: torch.nn.Module, device: str = "cuda") -> str: + """Print model summary using torchinfo (if available). + + Falls back to a simple parameter count if torchinfo is not installed. + + Returns: + Summary string. + """ + try: + from torchinfo import summary as torchinfo_summary + info = torchinfo_summary( + model, + input_data={ + "drone_img": torch.randn(1, 3, 256, 256, device=device), + "sat_img": torch.randn(1, 3, 256, 256, device=device), + }, + col_names=["input_size", "output_size", "num_params", "trainable"], + verbose=0, + depth=3, + ) + summary_str = str(info) + LOGGER.info("Model summary (torchinfo):\n%s", summary_str) + return summary_str + except ImportError: + LOGGER.info("torchinfo not installed, using basic parameter count") + except Exception as e: + LOGGER.warning("torchinfo failed (%s), using basic parameter count", e) + + # Fallback: simple param count. + lines = [] + total = 0 + trainable = 0 + for name, param in model.named_parameters(): + total += param.numel() + if param.requires_grad: + trainable += param.numel() + lines.append(f" [trainable] {name}: {list(param.shape)} ({param.numel():,})") + + summary_str = ( + f"Total parameters: {total:,}\n" + f"Trainable parameters: {trainable:,} ({100*trainable/max(total,1):.2f}%)\n" + + "\n".join(lines[:30]) + ) + LOGGER.info("Model summary:\n%s", summary_str) + return summary_str diff --git a/src/training/trackers.py b/src/training/trackers.py new file mode 100644 index 0000000..6673528 --- /dev/null +++ b/src/training/trackers.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +"""Unified experiment tracking: W&B + TensorBoard + CSV. + +Auto-detects available backends. Falls back gracefully if wandb/tensorboard +are not installed. + +Usage: + tracker = ExperimentTracker(output_dir, config_dict, use_wandb=True, use_tb=True) + tracker.log_train(epoch, {"loss": 0.5, "lr": 1e-4}) + tracker.log_val(epoch, {"r@1_q2g": 0.3}) + tracker.log_gradients(epoch, grad_norms_dict) + tracker.log_image(epoch, "gradcam/drone", image_tensor) + tracker.close() +""" + +import logging +from pathlib import Path +from typing import Any + +import torch + +LOGGER = logging.getLogger("caption_test.trackers") + + +def _try_import_wandb(): + try: + import wandb + return wandb + except ImportError: + return None + + +def _try_import_tb(): + try: + from torch.utils.tensorboard import SummaryWriter + return SummaryWriter + except ImportError: + return None + + +class ExperimentTracker: + """Unified tracker dispatching to W&B, TensorBoard, and CSV. + + Args: + output_dir: Base output directory. + config: Dict of hyperparameters to log. + use_wandb: Enable Weights & Biases tracking. + use_tb: Enable TensorBoard tracking. + wandb_project: W&B project name. + wandb_run_name: W&B run name (auto-generated if None). + wandb_entity: W&B entity (team/user). + """ + + def __init__( + self, + output_dir: str | Path, + config: dict[str, Any] | None = None, + use_wandb: bool = False, + use_tb: bool = True, + wandb_project: str = "caption-test-gtauav", + wandb_run_name: str | None = None, + wandb_entity: str | None = None, + ) -> None: + self.output_dir = Path(output_dir) + self._wandb_run = None + self._tb_writer = None + + # W&B init. + if use_wandb: + wandb = _try_import_wandb() + if wandb is not None: + self._wandb_run = wandb.init( + project=wandb_project, + name=wandb_run_name, + entity=wandb_entity, + config=config or {}, + dir=str(self.output_dir), + reinit=True, + ) + LOGGER.info("W&B initialized: %s", self._wandb_run.url) + else: + LOGGER.warning("wandb not installed, skipping W&B tracking") + + # TensorBoard init. + if use_tb: + SummaryWriter = _try_import_tb() + if SummaryWriter is not None: + tb_dir = self.output_dir / "tb_logs" + tb_dir.mkdir(parents=True, exist_ok=True) + self._tb_writer = SummaryWriter(log_dir=str(tb_dir)) + LOGGER.info("TensorBoard initialized: %s", tb_dir) + else: + LOGGER.warning("tensorboard not installed, skipping TB tracking") + + @property + def has_wandb(self) -> bool: + return self._wandb_run is not None + + @property + def has_tb(self) -> bool: + return self._tb_writer is not None + + def log_train(self, epoch: int, metrics: dict[str, float], step: int | None = None) -> None: + """Log training metrics for an epoch.""" + if self._wandb_run is not None: + self._wandb_run.log( + {f"train/{k}": v for k, v in metrics.items()}, + step=step or epoch, + ) + if self._tb_writer is not None: + for k, v in metrics.items(): + self._tb_writer.add_scalar(f"train/{k}", v, global_step=step or epoch) + + def log_val(self, epoch: int, metrics: dict[str, float], step: int | None = None) -> None: + """Log validation metrics.""" + if self._wandb_run is not None: + self._wandb_run.log( + {f"val/{k}": v for k, v in metrics.items()}, + step=step or epoch, + ) + if self._tb_writer is not None: + for k, v in metrics.items(): + self._tb_writer.add_scalar(f"val/{k}", v, global_step=step or epoch) + + def log_gradients(self, epoch: int, grad_norms: dict[str, float], step: int | None = None) -> None: + """Log gradient norms per parameter group.""" + if self._wandb_run is not None: + self._wandb_run.log( + {f"gradients/{k}": v for k, v in grad_norms.items()}, + step=step or epoch, + ) + if self._tb_writer is not None: + for k, v in grad_norms.items(): + self._tb_writer.add_scalar(f"gradients/{k}", v, global_step=step or epoch) + + def log_scalar(self, tag: str, value: float, step: int) -> None: + """Log a single scalar.""" + if self._wandb_run is not None: + self._wandb_run.log({tag: value}, step=step) + if self._tb_writer is not None: + self._tb_writer.add_scalar(tag, value, global_step=step) + + def log_image(self, tag: str, image: Any, step: int, caption: str | None = None) -> None: + """Log an image (numpy HWC or torch CHW). + + Args: + tag: Image tag/name. + image: numpy array [H,W,C] or torch tensor [C,H,W]. + step: Global step. + caption: Optional caption for W&B. + """ + if self._wandb_run is not None: + wandb = _try_import_wandb() + if isinstance(image, torch.Tensor): + image_np = image.detach().cpu().permute(1, 2, 0).numpy() + else: + image_np = image + self._wandb_run.log( + {tag: wandb.Image(image_np, caption=caption)}, + step=step, + ) + if self._tb_writer is not None: + if isinstance(image, torch.Tensor): + self._tb_writer.add_image(tag, image.detach().cpu(), global_step=step) + else: + self._tb_writer.add_image(tag, image, global_step=step, dataformats="HWC") + + def log_histogram(self, tag: str, values: torch.Tensor, step: int) -> None: + """Log a histogram of values (weights, activations, etc.).""" + if self._wandb_run is not None: + wandb = _try_import_wandb() + self._wandb_run.log( + {tag: wandb.Histogram(values.detach().cpu().numpy())}, + step=step, + ) + if self._tb_writer is not None: + self._tb_writer.add_histogram(tag, values.detach().cpu(), global_step=step) + + def log_model_graph(self, model: torch.nn.Module, input_example: Any = None) -> None: + """Log model graph to TensorBoard (if available).""" + if self._tb_writer is not None and input_example is not None: + try: + self._tb_writer.add_graph(model, input_example) + except Exception as e: + LOGGER.warning("Failed to log model graph: %s", e) + + def watch_model(self, model: torch.nn.Module, log_freq: int = 100) -> None: + """Enable W&B gradient/weight watching.""" + if self._wandb_run is not None: + wandb = _try_import_wandb() + wandb.watch(model, log="all", log_freq=log_freq) + + def log_summary(self, summary: dict[str, Any]) -> None: + """Log final summary metrics (best R@1, etc.).""" + if self._wandb_run is not None: + for k, v in summary.items(): + self._wandb_run.summary[k] = v + + def close(self) -> None: + """Flush and close all backends.""" + if self._tb_writer is not None: + self._tb_writer.flush() + self._tb_writer.close() + if self._wandb_run is not None: + self._wandb_run.finish() diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index c2736b5..5a1c55e 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -4,6 +4,9 @@ from __future__ import annotations Asymmetric DINOv3 encoders (drone LVD + satellite SAT) with LRSCLIP text fusion. Single InfoNCE loss: query(drone+text) vs gallery(satellite). + +Supports gin-config, W&B, TensorBoard, Grad-CAM, gradient monitoring, +PyTorch Profiler, and torchinfo model summary. """ import argparse @@ -16,6 +19,7 @@ from dataclasses import dataclass, field from pathlib import Path import coloredlogs +import gin import pandas as pd import torch import torch.nn as nn @@ -28,6 +32,9 @@ from tqdm import tqdm from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch from src.losses.multi_infonce import InfoNCELoss 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 +from src.training.profiling import TrainingProfiler, print_model_summary from src.models.asymmetric_encoder import ( AsymmetricEncoder, get_dino_transform, @@ -48,6 +55,7 @@ _DINO_SAT = "nn_models/DINO_SAT/model.safetensors" _LRSCLIP = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt" +@gin.configurable @dataclass class TrainConfigGTAUAV: """Training configuration for GTA-UAV experiment.""" @@ -69,7 +77,7 @@ class TrainConfigGTAUAV: # Training. resume_from: str | None = None # path to checkpoint for resuming output_dir: str = "out/gtauav/with_text" - epochs: int = 20 + epochs: int = 10 batch_size: int = 8 num_workers: int = 4 learning_rate: float = 1e-4 @@ -89,6 +97,20 @@ class TrainConfigGTAUAV: weight_g2q: float = 0.4 learnable_temperature: bool = True + # Tracking & diagnostics. + use_wandb: bool = False + use_tb: bool = True + wandb_project: str = "caption-test-gtauav" + wandb_run_name: str | None = None + wandb_entity: str | None = None + log_grad_norms: bool = True + use_gradcam: bool = False + gradcam_every: int = 5 # Grad-CAM every N epochs + gradcam_samples: int = 8 + use_profiler: bool = False + profiler_warmup: int = 3 + profiler_active: int = 5 + def _set_seed(seed: int) -> None: import random as _random @@ -157,7 +179,7 @@ def _evaluate( all_query: list[torch.Tensor] = [] all_gallery: list[torch.Tensor] = [] - for batch in tqdm(loader, desc=" 🔎 Evaluating", unit="batch", leave=False): + for batch in tqdm(loader, desc=" eval", unit="batch", leave=False): drone_img = batch["drone_img"].to(device, non_blocking=True) sat_img = batch["sat_img"].to(device, non_blocking=True) @@ -240,7 +262,7 @@ def _clear_vram() -> None: torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() allocated = torch.cuda.memory_allocated() / 1e9 - LOGGER.info("🧹 VRAM cleared. Current usage: %.2f GB", allocated) + LOGGER.info("VRAM cleared. Current usage: %.2f GB", allocated) def train(cfg: TrainConfigGTAUAV) -> None: @@ -259,12 +281,23 @@ def train(cfg: TrainConfigGTAUAV) -> None: with (output_dir / "config.json").open("w") as f: json.dump(vars(cfg), f, indent=2) + # --- Experiment tracker (W&B + TensorBoard) --- + tracker = ExperimentTracker( + output_dir=output_dir, + config=vars(cfg), + use_wandb=cfg.use_wandb, + use_tb=cfg.use_tb, + wandb_project=cfg.wandb_project, + wandb_run_name=cfg.wandb_run_name, + wandb_entity=cfg.wandb_entity, + ) + # Model. start_epoch = 0 resume_ckpt = None if cfg.resume_from is not None: - LOGGER.info("🔄 Resuming from %s", cfg.resume_from) + LOGGER.info("Resuming from %s", cfg.resume_from) model, resume_ckpt = AsymmetricEncoder.load_checkpoint( cfg.resume_from, dino_web_path=cfg.dino_web_path, @@ -274,8 +307,8 @@ def train(cfg: TrainConfigGTAUAV) -> None: ) start_epoch = resume_ckpt.get("epoch", -1) + 1 else: - mode_str = "🚫 baseline (no text)" if cfg.baseline_mode else "📝 with text (L1/L2/L3)" - LOGGER.info("🏗️ Building model — %s", mode_str) + mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)" + LOGGER.info("Building model — %s", mode_str) model = AsymmetricEncoder( dino_web_path=cfg.dino_web_path, dino_sat_path=cfg.dino_sat_path, @@ -288,10 +321,18 @@ def train(cfg: TrainConfigGTAUAV) -> None: n_trainable = sum(p.numel() for p in model.trainable_parameters()) n_total = sum(p.numel() for p in model.parameters()) LOGGER.info( - "🧮 trainable=%s (%.2f%%) total=%s", + "trainable=%s (%.2f%%) total=%s", f"{n_trainable:,}", 100.0 * n_trainable / max(n_total, 1), f"{n_total:,}", ) + # --- Model summary (torchinfo) --- + model_summary = print_model_summary(model, device=cfg.device) + (output_dir / "model_summary.txt").write_text(model_summary) + + # --- W&B model watching (gradient + weight histograms) --- + if tracker.has_wandb: + tracker.watch_model(model, log_freq=50) + # Loss. loss_fn = InfoNCELoss( temperature_init=cfg.tau_init, @@ -301,7 +342,7 @@ def train(cfg: TrainConfigGTAUAV) -> None: learnable_temperature=cfg.learnable_temperature, ) LOGGER.info( - "🌡️ Temperature: %s (init=%.3f)", + "Temperature: %s (init=%.3f)", "learnable" if cfg.learnable_temperature else "cosine schedule", cfg.tau_init, ) @@ -345,7 +386,7 @@ def train(cfg: TrainConfigGTAUAV) -> None: pin_memory=True, ) - LOGGER.info("📦 train=%d test=%d batch=%d", len(train_ds), len(test_ds), cfg.batch_size) + LOGGER.info("train=%d test=%d batch=%d", len(train_ds), len(test_ds), cfg.batch_size) # Optimizer — per-group LR (text encoder gets lower LR). param_groups = _build_param_groups(model, cfg.learning_rate, cfg.text_lr_factor) @@ -358,7 +399,7 @@ def train(cfg: TrainConfigGTAUAV) -> None: lr_info = f"proj={cfg.learning_rate:.0e}" if not cfg.baseline_mode: lr_info += f" text={cfg.learning_rate * cfg.text_lr_factor:.0e}" - LOGGER.info("⚙️ Optimizer: AdamW LR: %s warmup=%d epochs", lr_info, cfg.warmup_epochs) + LOGGER.info("Optimizer: AdamW LR: %s warmup=%d epochs", lr_info, cfg.warmup_epochs) # Scheduler — cosine with linear warmup. steps_per_epoch = len(train_loader) @@ -377,18 +418,31 @@ def train(cfg: TrainConfigGTAUAV) -> None: if resume_ckpt is not None: if "optimizer_state" in resume_ckpt: optimizer.load_state_dict(resume_ckpt["optimizer_state"]) - LOGGER.info("🔄 Optimizer state restored") + LOGGER.info("Optimizer state restored") if "loss_state" in resume_ckpt: loss_fn.load_state_dict(resume_ckpt["loss_state"]) - LOGGER.info("🔄 Loss state restored (tau=%.4f)", loss_fn.current_temperature) + LOGGER.info("Loss state restored (tau=%.4f)", loss_fn.current_temperature) # Set scheduler last_epoch so it resumes at the correct LR. scheduler.last_epoch = start_epoch * steps_per_epoch - LOGGER.info("🔄 Resuming from epoch %d", start_epoch) + LOGGER.info("Resuming from epoch %d", start_epoch) history: list[dict] = [] csv_logger = CSVLogger(output_dir) - LOGGER.info("🚀 Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch) + # --- Optional profiler (first epoch only) --- + profiler = None + if cfg.use_profiler and start_epoch == 0: + profiler = TrainingProfiler( + output_dir=output_dir, + n_warmup=cfg.profiler_warmup, + n_active=cfg.profiler_active, + ) + profiler.start() + + LOGGER.info("Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch) + + global_step = start_epoch * steps_per_epoch + best_r1 = 0.0 for epoch in range(start_epoch, cfg.epochs): model.train() @@ -398,7 +452,7 @@ def train(cfg: TrainConfigGTAUAV) -> None: pbar = tqdm( train_loader, - desc=f" 🏋️ Epoch {epoch + 1}/{cfg.epochs}", + desc=f" Epoch {epoch + 1}/{cfg.epochs}", unit="batch", leave=False, ) @@ -439,15 +493,34 @@ def train(cfg: TrainConfigGTAUAV) -> None: model.trainable_parameters(), max_norm=cfg.grad_clip, ) + + # --- Gradient monitoring (after unscale, before step) --- + if cfg.log_grad_norms and n_batches % 50 == 0: + grad_norms = compute_gradient_norms(model, loss_fn) + tracker.log_gradients(epoch, grad_norms, step=global_step) + if n_batches == 0: + log_gradient_summary(grad_norms) + scaler.step(optimizer) scaler.update() with warnings.catch_warnings(): warnings.filterwarnings("ignore", message=".*lr_scheduler.step.*optimizer.step.*") scheduler.step() + # --- Per-step tracking --- + step_metrics = { + "loss": float(total_loss.item()), + "temperature": float(loss_dict["temperature"].item()), + "gate_q": float(loss_dict["gate_q"].item()), + "gate_g": float(loss_dict["gate_g"].item()), + "lr": optimizer.param_groups[0]["lr"], + } + tracker.log_train(epoch, step_metrics, step=global_step) + for key, val in loss_dict.items(): agg[key] = agg.get(key, 0.0) + float(val.item()) n_batches += 1 + global_step += 1 pbar.set_postfix( loss=f"{total_loss.item():.3f}", @@ -456,11 +529,18 @@ def train(cfg: TrainConfigGTAUAV) -> None: gg=f"{loss_dict['gate_g'].item():.3f}", ) + # --- Profiler step --- + if profiler is not None: + profiler.step() + if profiler.is_done(n_batches): + profiler.export() + profiler = None + elapsed = time.time() - epoch_start means = {k: v / max(n_batches, 1) for k, v in agg.items()} LOGGER.info( - "📈 epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate_q=%.4f gate_g=%.4f", + "epoch=%d time=%.1fs lr=%.2e loss=%.4f tau=%.4f gate_q=%.4f gate_g=%.4f", epoch, elapsed, optimizer.param_groups[0]["lr"], means.get("total", 0.0), @@ -475,8 +555,14 @@ def train(cfg: TrainConfigGTAUAV) -> None: "train": means, } - # Log train metrics to CSV. + # Log train metrics to CSV + generate plots every epoch. csv_logger.log_train(epoch, means, optimizer.param_groups[0]["lr"], elapsed) + generate_plots(csv_logger.log_dir) + + # --- Log VRAM usage --- + if torch.cuda.is_available(): + vram_gb = torch.cuda.max_memory_allocated() / 1e9 + tracker.log_scalar("system/vram_peak_gb", vram_gb, step=global_step) # Evaluation. if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1: @@ -484,8 +570,16 @@ def train(cfg: TrainConfigGTAUAV) -> None: epoch_record["val"] = val_metrics csv_logger.log_val(epoch, val_metrics) generate_plots(csv_logger.log_dir) + tracker.log_val(epoch, val_metrics, step=global_step) + + # Track best R@1. + r1 = val_metrics.get("r@1_q2g", 0.0) + if r1 > best_r1: + best_r1 = r1 + tracker.log_scalar("val/best_r@1_q2g", best_r1, step=global_step) + LOGGER.info( - "🎯 val epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f", + "val epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f", epoch, val_metrics.get("r@1_q2g", 0.0), val_metrics.get("r@5_q2g", 0.0), @@ -494,6 +588,27 @@ def train(cfg: TrainConfigGTAUAV) -> None: val_metrics.get("gate_g", 1.0), ) + # --- Grad-CAM visualization --- + if cfg.use_gradcam and (epoch + 1) % cfg.gradcam_every == 0: + from src.training.gradcam import generate_gradcam_samples + overlays = generate_gradcam_samples( + model=model, + dataloader=test_loader, + device=cfg.device, + output_dir=str(output_dir), + n_samples=cfg.gradcam_samples, + epoch=epoch, + ) + # Log first few overlays to tracker. + for i, overlay in enumerate(overlays[:4]): + kind = "drone" if i % 2 == 0 else "sat" + tracker.log_image( + f"gradcam/{kind}_{i//2}", + overlay, + step=global_step, + caption=f"Epoch {epoch} {kind} Grad-CAM", + ) + history.append(epoch_record) # Save checkpoint. @@ -506,7 +621,7 @@ def train(cfg: TrainConfigGTAUAV) -> None: }, path=output_dir / f"ckpt_epoch{epoch:03d}.pt", ) - LOGGER.info("💾 Checkpoint saved: ckpt_epoch%03d.pt", epoch) + LOGGER.info("Checkpoint saved: ckpt_epoch%03d.pt", epoch) # Save history. history_path = output_dir / "history.json" @@ -514,7 +629,7 @@ def train(cfg: TrainConfigGTAUAV) -> None: json.dump(history, f, indent=2) # Save final eval report. - LOGGER.info("🔎 Running final evaluation...") + LOGGER.info("Running final evaluation...") final_metrics = _evaluate(model, test_loader, cfg.device) report = { "config": vars(cfg), @@ -525,9 +640,25 @@ def train(cfg: TrainConfigGTAUAV) -> None: with report_path.open("w", encoding="utf-8") as f: json.dump(report, f, indent=2) - LOGGER.info("✅ Training complete. Report: %s", report_path) + # --- Log final summary to W&B --- + tracker.log_summary({ + "best_r@1_q2g": best_r1, + "final_r@1_q2g": final_metrics.get("r@1_q2g", 0.0), + "final_r@5_q2g": final_metrics.get("r@5_q2g", 0.0), + "final_r@10_q2g": final_metrics.get("r@10_q2g", 0.0), + "final_gate_q": final_metrics.get("gate_q", 1.0), + "final_gate_g": final_metrics.get("gate_g", 1.0), + }) + + # --- Cleanup profiler if still running --- + if profiler is not None: + profiler.export() + + tracker.close() + + LOGGER.info("Training complete. Report: %s", report_path) LOGGER.info( - "📊 Final — R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f", + "Final — R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f", final_metrics.get("r@1_q2g", 0.0), final_metrics.get("r@5_q2g", 0.0), final_metrics.get("r@10_q2g", 0.0), @@ -538,6 +669,10 @@ def train(cfg: TrainConfigGTAUAV) -> None: def main() -> None: parser = argparse.ArgumentParser(description="GTA-UAV caption test training.") + parser.add_argument( + "--config", type=str, default=None, + help="Path to gin config file (e.g. conf/gtauav_balanced.gin).", + ) parser.add_argument( "--baseline", action="store_true", help="Run baseline mode (no text).", @@ -555,47 +690,86 @@ def main() -> None: help="Path to seg_filter.json for excluding bad images.", ) parser.add_argument( - "--batch-size", type=int, default=8, + "--batch-size", type=int, default=None, help="Batch size.", ) parser.add_argument( - "--epochs", type=int, default=20, + "--epochs", type=int, default=None, help="Number of epochs.", ) parser.add_argument( - "--lr", type=float, default=1e-4, + "--lr", type=float, default=None, help="Learning rate for projections.", ) parser.add_argument( - "--text-lr-factor", type=float, default=0.1, + "--text-lr-factor", type=float, default=None, help="Text encoder LR = lr * factor (default 0.1 = 10x lower).", ) parser.add_argument( - "--warmup-epochs", type=int, default=2, + "--warmup-epochs", type=int, default=None, help="Linear warmup epochs.", ) parser.add_argument( - "--init-gate", type=float, default=0.7, + "--init-gate", type=float, default=None, help="Initial gate value (image weight).", ) + # Tracking flags. + parser.add_argument("--wandb", action="store_true", help="Enable W&B tracking.") + parser.add_argument("--no-tb", action="store_true", help="Disable TensorBoard.") + parser.add_argument("--gradcam", action="store_true", help="Enable Grad-CAM visualization.") + parser.add_argument("--profile", action="store_true", help="Enable PyTorch profiler (first epoch).") + parser.add_argument("--no-grad-norms", action="store_true", help="Disable gradient norm logging.") + # Gin overrides. + parser.add_argument( + "--gin-param", type=str, nargs="*", default=[], + help="Gin parameter overrides (e.g. 'TrainConfigGTAUAV.epochs=30').", + ) args = parser.parse_args() - cfg = TrainConfigGTAUAV() - cfg.baseline_mode = args.baseline - cfg.resume_from = args.resume - cfg.batch_size = args.batch_size - cfg.epochs = args.epochs - cfg.learning_rate = args.lr - cfg.text_lr_factor = args.text_lr_factor - cfg.warmup_epochs = args.warmup_epochs - cfg.init_gate = args.init_gate + # Parse gin config if provided. + if args.config is not None: + gin.parse_config_file(args.config) + if args.gin_param: + gin.parse_config(args.gin_param) + # Create config (gin bindings apply via @gin.configurable). + cfg = TrainConfigGTAUAV() + + # CLI overrides take priority over gin. + if args.baseline: + cfg.baseline_mode = True + if args.resume is not None: + cfg.resume_from = args.resume + if args.batch_size is not None: + cfg.batch_size = args.batch_size + if args.epochs is not None: + cfg.epochs = args.epochs + if args.lr is not None: + cfg.learning_rate = args.lr + if args.text_lr_factor is not None: + cfg.text_lr_factor = args.text_lr_factor + if args.warmup_epochs is not None: + cfg.warmup_epochs = args.warmup_epochs + if args.init_gate is not None: + cfg.init_gate = args.init_gate if args.filter_meta is not None: cfg.filter_meta = args.filter_meta + # Tracking overrides. + if args.wandb: + cfg.use_wandb = True + if args.no_tb: + cfg.use_tb = False + if args.gradcam: + cfg.use_gradcam = True + if args.profile: + cfg.use_profiler = True + if args.no_grad_norms: + cfg.log_grad_norms = False + if args.output_dir is not None: cfg.output_dir = args.output_dir - elif args.baseline: + elif args.baseline and args.output_dir is None: cfg.output_dir = "out/gtauav/baseline" train(cfg)