Improve training: learnable temperature, per-group LR, warmup, augmentations

Loss:
- Learnable temperature (CLIP-style logit_scale) with clamp [0.01, 0.5]
- Replaces fixed cosine schedule (still available via --no-learnable-temp)
- Default tau_init=0.07

Optimizer:
- Per-group LR: projections 1e-4, text encoder 1e-5 (10x lower)
- Learnable temperature included in projection param group

Scheduler:
- Linear warmup (2 epochs default) + cosine annealing
- Per-step scheduling (not per-epoch)

Augmentations (separate drone/satellite):
- Drone: RandomResizedCrop(0.7-1.0), HFlip, Rotation(15), ColorJitter,
  RandomGrayscale(0.05), GaussianBlur
- Satellite: RandomResizedCrop(0.7-1.0), HFlip, ColorJitter, RandomGrayscale
- Eval: clean Resize+CenterCrop (no augmentation)

Dataset: supports separate drone_transform/sat_transform args

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

src/datasets/gtauav_dataset.py

@@ -99,7 +99,9 @@ class GTAUAVDataset(Dataset):
         pair_json: Path to cross-area-drone2sate-{train,test}.json.
         rgb_root: Root of GTA-UAV-LR RGB images.
         caption_root: Root of GTA-UAV-LR-captions.
-        image_transform: Callable applied to PIL images.
+        drone_transform: Transform for drone images (can include augmentations).
+        sat_transform: Transform for satellite images (can include augmentations).
+        image_transform: Fallback single transform for both (used if drone/sat not set).
         filter_meta: Path to seg_filter.json (exclude 90%+ bg/water).
         drop_caption_prob: Probability of dropping captions (ablation).
         seed: Random seed.
@@ -110,6 +112,8 @@ class GTAUAVDataset(Dataset):
         pair_json: str,
         rgb_root: str = str(_RGB_ROOT),
         caption_root: str = str(_CAPTION_ROOT),
+        drone_transform: Callable[[Image.Image], torch.Tensor] | None = None,
+        sat_transform: Callable[[Image.Image], torch.Tensor] | None = None,
         image_transform: Callable[[Image.Image], torch.Tensor] | None = None,
         filter_meta: str | None = None,
         drop_caption_prob: float = 0.0,
@@ -117,7 +121,8 @@ class GTAUAVDataset(Dataset):
     ) -> None:
         self.rgb_root = Path(rgb_root)
         self.caption_root = Path(caption_root)
-        self.image_transform = image_transform
+        self.drone_transform = drone_transform or image_transform
+        self.sat_transform = sat_transform or image_transform
         self.drop_caption_prob = drop_caption_prob
         self._rng = random.Random(seed)
 
@@ -189,12 +194,12 @@ class GTAUAVDataset(Dataset):
                 "caption_l3": l3,
             })
 
-    def _load_image(self, directory: str, filename: str) -> torch.Tensor:
+    def _load_image(self, directory: str, filename: str, transform: Callable | None = None) -> torch.Tensor:
         path = self.rgb_root / directory / filename
         with Image.open(path) as img:
             rgb = img.convert("RGB")
-        if self.image_transform is not None:
+        if transform is not None:
-            return self.image_transform(rgb)
+            return transform(rgb)
         return torch.tensor(0)  # placeholder if no transform
 
     def __len__(self) -> int:
@@ -203,7 +208,7 @@ class GTAUAVDataset(Dataset):
     def __getitem__(self, idx: int) -> dict[str, Any]:
         entry = self.entries[idx]
 
-        drone_img = self._load_image(entry["drone_dir"], entry["drone_name"])
+        drone_img = self._load_image(entry["drone_dir"], entry["drone_name"], self.drone_transform)
 
         # Sample satellite match (weighted if semi-positive).
         if entry["sat_weights"] is not None:
@@ -215,7 +220,7 @@ class GTAUAVDataset(Dataset):
         else:
             sat_name = self._rng.choice(entry["sat_candidates"])
 
-        sat_img = self._load_image(entry["sat_dir"], sat_name)
+        sat_img = self._load_image(entry["sat_dir"], sat_name, self.sat_transform)
 
         # Captions with optional dropout.
         if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:

src/losses/multi_infonce.py

@@ -4,7 +4,8 @@ from __future__ import annotations
 
 Single symmetric InfoNCE between query (drone+text fused) and gallery (satellite).
 Asymmetric weighting: query->gallery weighted higher (real use-case direction).
-Cosine temperature schedule for sharper distribution over training.
+
+Supports both learnable temperature (CLIP-style logit_scale) and fixed/scheduled.
 """
 
 import math
@@ -47,23 +48,30 @@ def cosine_temperature(
 
 @gin.configurable
 class InfoNCELoss(nn.Module):
-    """Symmetric InfoNCE with cosine temperature schedule.
+    """Symmetric InfoNCE with learnable or scheduled temperature.
 
     Args:
-        temperature_init: Temperature at epoch 0.
+        temperature_init: Initial temperature value.
-        temperature_final: Temperature after cosine decay.
+        temperature_final: Final temperature (only used if learnable=False).
         label_smoothing: Cross-entropy label smoothing.
         weight_q2g: Weight for query->gallery direction.
         weight_g2q: Weight for gallery->query direction.
+        learnable_temperature: If True, temperature is a learnable parameter
+            (CLIP-style logit_scale). If False, uses cosine schedule.
+        tau_min: Minimum clamp for learnable temperature.
+        tau_max: Maximum clamp for learnable temperature.
     """
 
     def __init__(
         self,
-        temperature_init: float = 0.1,
+        temperature_init: float = 0.07,
         temperature_final: float = 0.01,
         label_smoothing: float = 0.1,
         weight_q2g: float = 0.6,
         weight_g2q: float = 0.4,
+        learnable_temperature: bool = True,
+        tau_min: float = 0.01,
+        tau_max: float = 0.5,
     ) -> None:
         super().__init__()
         self.temperature_init = temperature_init
@@ -71,6 +79,26 @@ class InfoNCELoss(nn.Module):
         self.label_smoothing = label_smoothing
         self.weight_q2g = weight_q2g
         self.weight_g2q = weight_g2q
+        self.learnable_temperature = learnable_temperature
+        self.tau_min = tau_min
+        self.tau_max = tau_max
+
+        if learnable_temperature:
+            # Store as log(1/tau) like CLIP's logit_scale.
+            init_logit_scale = math.log(1.0 / temperature_init)
+            self.logit_scale = nn.Parameter(torch.tensor(init_logit_scale))
+        else:
+            self.logit_scale = None
+
+    @property
+    def current_temperature(self) -> float:
+        """Current temperature value (for logging)."""
+        if self.logit_scale is not None:
+            tau = 1.0 / self.logit_scale.exp().clamp(
+                min=1.0 / self.tau_max, max=1.0 / self.tau_min,
+            ).item()
+            return tau
+        return self.temperature_init
 
     def forward(
         self,
@@ -89,6 +117,13 @@ class InfoNCELoss(nn.Module):
         Returns:
             Dict with 'total', 'temperature', 'gate'.
         """
+        if self.learnable_temperature:
+            # Clamp logit_scale to prevent tau from going out of bounds.
+            logit_scale = self.logit_scale.exp().clamp(
+                min=1.0 / self.tau_max, max=1.0 / self.tau_min,
+            )
+            tau = 1.0 / logit_scale
+        else:
             tau = cosine_temperature(
                 epoch=epoch,
                 total_epochs=total_epochs,
@@ -107,8 +142,13 @@ class InfoNCELoss(nn.Module):
 
         gate = embeddings.get("gate", 1.0)
 
+        if isinstance(tau, float):
+            tau_out = torch.tensor(tau, device=loss.device)
+        else:
+            tau_out = tau.detach().clone()
+
         return {
             "total": loss,
-            "temperature": torch.tensor(tau, device=loss.device),
+            "temperature": tau_out,
             "gate": torch.tensor(gate, device=loss.device),
         }

src/models/asymmetric_encoder.py

@@ -543,15 +543,58 @@ class AsymmetricEncoder(nn.Module):
 # Image preprocessing (DINOv3: 256x256, ImageNet normalization)
 # ---------------------------------------------------------------------------
 
+_IMAGENET_MEAN = [0.485, 0.456, 0.406]
+_IMAGENET_STD = [0.229, 0.224, 0.225]
+
+
 def get_dino_transform(image_size: int = 256) -> torch.nn.Module:
-    """Build image transform for DINOv3 input."""
+    """Build eval/inference image transform for DINOv3 input."""
     from torchvision import transforms
     return transforms.Compose([
         transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BICUBIC),
         transforms.CenterCrop(image_size),
         transforms.ToTensor(),
-        transforms.Normalize(
+        transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
-            mean=[0.485, 0.456, 0.406],
+    ])
-            std=[0.229, 0.224, 0.225],
+
-        ),
+
+def get_drone_train_transform(image_size: int = 256) -> torch.nn.Module:
+    """Build training augmentation for drone images.
+
+    Includes: RandomResizedCrop, HFlip, rotation, color jitter,
+    grayscale, Gaussian blur.
+    """
+    from torchvision import transforms
+    return transforms.Compose([
+        transforms.RandomResizedCrop(
+            image_size, scale=(0.7, 1.0),
+            interpolation=transforms.InterpolationMode.BICUBIC,
+        ),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.RandomRotation(degrees=15),
+        transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1),
+        transforms.RandomGrayscale(p=0.05),
+        transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
+    ])
+
+
+def get_satellite_train_transform(image_size: int = 256) -> torch.nn.Module:
+    """Build training augmentation for satellite images.
+
+    Lighter than drone: no rotation or blur (satellite is nadir/consistent).
+    Includes: RandomResizedCrop, HFlip, color jitter, grayscale.
+    """
+    from torchvision import transforms
+    return transforms.Compose([
+        transforms.RandomResizedCrop(
+            image_size, scale=(0.7, 1.0),
+            interpolation=transforms.InterpolationMode.BICUBIC,
+        ),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1),
+        transforms.RandomGrayscale(p=0.05),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
     ])

src/training/train_gtauav.py

@@ -9,6 +9,7 @@ Single InfoNCE loss: query(drone+text) vs gallery(satellite).
 import argparse
 import json
 import logging
+import math
 import time
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -18,13 +19,18 @@ import torch
 import torch.nn as nn
 from torch.amp import GradScaler, autocast
 from torch.optim import AdamW
-from torch.optim.lr_scheduler import CosineAnnealingLR
+from torch.optim.lr_scheduler import LambdaLR
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
 from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch
 from src.losses.multi_infonce import InfoNCELoss
-from src.models.asymmetric_encoder import AsymmetricEncoder, get_dino_transform
+from src.models.asymmetric_encoder import (
+    AsymmetricEncoder,
+    get_dino_transform,
+    get_drone_train_transform,
+    get_satellite_train_transform,
+)
 
 LOGGER = logging.getLogger("caption_test.train_gtauav")
 
@@ -64,19 +70,21 @@ class TrainConfigGTAUAV:
     batch_size: int = 64
     num_workers: int = 4
     learning_rate: float = 1e-4
+    text_lr_factor: float = 0.1  # text encoder LR = learning_rate * factor
     weight_decay: float = 1e-4
     grad_clip: float = 1.0
     use_amp: bool = True
     eval_every: int = 2
+    warmup_epochs: int = 2
     seed: int = 42
     device: str = "cuda"
 
     # Loss.
-    tau_init: float = 0.1
+    tau_init: float = 0.07
-    tau_final: float = 0.01
     label_smoothing: float = 0.1
     weight_q2g: float = 0.6
     weight_g2q: float = 0.4
+    learnable_temperature: bool = True
 
 
 def _set_seed(seed: int) -> None:
@@ -95,6 +103,45 @@ def _atomic_save(obj: dict, path: Path) -> None:
     tmp_path.replace(path)
 
 
+def _build_param_groups(
+    model: AsymmetricEncoder,
+    lr: float,
+    text_lr_factor: float,
+) -> list[dict]:
+    """Build optimizer param groups with separate LR for text encoder."""
+    text_params = []
+    other_params = []
+
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        if "text_encoder" in name:
+            text_params.append(param)
+        else:
+            other_params.append(param)
+
+    groups = [{"params": other_params, "lr": lr}]
+    if text_params:
+        groups.append({"params": text_params, "lr": lr * text_lr_factor})
+
+    return groups
+
+
+def _cosine_warmup_schedule(
+    warmup_steps: int,
+    total_steps: int,
+) -> callable:
+    """Cosine annealing with linear warmup."""
+
+    def lr_lambda(step: int) -> float:
+        if step < warmup_steps:
+            return step / max(warmup_steps, 1)
+        progress = (step - warmup_steps) / max(total_steps - warmup_steps, 1)
+        return 0.5 * (1.0 + math.cos(math.pi * progress))
+
+    return lr_lambda
+
+
 @torch.no_grad()
 def _evaluate(
     model: AsymmetricEncoder,
@@ -186,27 +233,35 @@ def train(cfg: TrainConfigGTAUAV) -> None:
     # Loss.
     loss_fn = InfoNCELoss(
         temperature_init=cfg.tau_init,
-        temperature_final=cfg.tau_final,
         label_smoothing=cfg.label_smoothing,
         weight_q2g=cfg.weight_q2g,
         weight_g2q=cfg.weight_g2q,
+        learnable_temperature=cfg.learnable_temperature,
+    )
+    LOGGER.info(
+        "🌡️  Temperature: %s (init=%.3f)",
+        "learnable" if cfg.learnable_temperature else "cosine schedule",
+        cfg.tau_init,
     )
 
-    # Data.
+    # Data — separate transforms for train (augmented) and eval (clean).
-    transform = get_dino_transform(image_size=256)
+    drone_train_tf = get_drone_train_transform(image_size=256)
+    sat_train_tf = get_satellite_train_transform(image_size=256)
+    eval_tf = get_dino_transform(image_size=256)
 
     train_ds = GTAUAVDataset(
         pair_json=cfg.train_json,
         rgb_root=cfg.rgb_root,
         caption_root=cfg.caption_root,
-        image_transform=transform,
+        drone_transform=drone_train_tf,
+        sat_transform=sat_train_tf,
         filter_meta=cfg.filter_meta,
     )
     test_ds = GTAUAVDataset(
         pair_json=cfg.test_json,
         rgb_root=cfg.rgb_root,
         caption_root=cfg.caption_root,
-        image_transform=transform,
+        image_transform=eval_tf,
         filter_meta=cfg.filter_meta,
     )
 
@@ -230,13 +285,27 @@ def train(cfg: TrainConfigGTAUAV) -> None:
 
     LOGGER.info("📦 train=%d  test=%d  batch=%d", len(train_ds), len(test_ds), cfg.batch_size)
 
-    # Optimizer.
+    # Optimizer — per-group LR (text encoder gets lower LR).
-    optimizer = AdamW(
+    param_groups = _build_param_groups(model, cfg.learning_rate, cfg.text_lr_factor)
-        model.trainable_parameters(),
+    # Include loss temperature if learnable.
-        lr=cfg.learning_rate,
+    if cfg.learnable_temperature and loss_fn.logit_scale is not None:
-        weight_decay=cfg.weight_decay,
+        param_groups[0]["params"].append(loss_fn.logit_scale)
+
+    optimizer = AdamW(param_groups, weight_decay=cfg.weight_decay)
+
+    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)
+
+    # Scheduler — cosine with linear warmup.
+    steps_per_epoch = len(train_loader)
+    total_steps = cfg.epochs * steps_per_epoch
+    warmup_steps = cfg.warmup_epochs * steps_per_epoch
+    scheduler = LambdaLR(
+        optimizer,
+        lr_lambda=_cosine_warmup_schedule(warmup_steps, total_steps),
     )
-    scheduler = CosineAnnealingLR(optimizer, T_max=cfg.epochs)
     scaler = GradScaler(enabled=cfg.use_amp)
 
     history: list[dict] = []
@@ -289,6 +358,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
                 )
             scaler.step(optimizer)
             scaler.update()
+            scheduler.step()
 
             for key, val in loss_dict.items():
                 agg[key] = agg.get(key, 0.0) + float(val.item())
@@ -296,10 +366,10 @@ def train(cfg: TrainConfigGTAUAV) -> None:
 
             pbar.set_postfix(
                 loss=f"{total_loss.item():.3f}",
+                tau=f"{loss_dict['temperature'].item():.4f}",
                 gate=f"{loss_dict['gate'].item():.3f}",
             )
 
-        scheduler.step()
         elapsed = time.time() - epoch_start
 
         means = {k: v / max(n_batches, 1) for k, v in agg.items()}
@@ -339,6 +409,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
                 "epoch": epoch,
                 "model_state": model.state_dict(),
                 "optimizer_state": optimizer.state_dict(),
+                "loss_state": loss_fn.state_dict(),
             },
             path=output_dir / f"ckpt_epoch{epoch:03d}.pt",
         )
@@ -395,7 +466,15 @@ def main() -> None:
     )
     parser.add_argument(
         "--lr", type=float, default=1e-4,
-        help="Learning rate.",
+        help="Learning rate for projections.",
+    )
+    parser.add_argument(
+        "--text-lr-factor", type=float, default=0.1,
+        help="Text encoder LR = lr * factor (default 0.1 = 10x lower).",
+    )
+    parser.add_argument(
+        "--warmup-epochs", type=int, default=2,
+        help="Linear warmup epochs.",
     )
     parser.add_argument(
         "--init-gate", type=float, default=0.7,
@@ -408,6 +487,8 @@ def main() -> None:
     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
 
     if args.filter_meta is not None: