Add ML diagnostics tooling (W&B, TensorBoard, Grad-CAM, profiler) and gin configs

- Add unified experiment tracker (W&B + TensorBoard) with graceful fallback
- Add gradient norm monitoring per param group (MONA, LoRA, MLP, gates, tau)
- Add Grad-CAM visualization for DINOv3 drone/satellite encoders
- Add PyTorch Profiler wrapper + torchinfo model summary
- Add gin-config support to train_gtauav.py with CLI overrides
- Add v3 gin configs: gtauav_balanced, gtauav_baseline, gtauav_text_heavy, gtauav_image_heavy
- Generate metric plots every epoch (not just on eval)
- Set default epochs to 10
- Update README and CLAUDE.md with new tooling and usage docs

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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

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

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

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

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)