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

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

Browse Source 
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>
This commit is contained in:
pikaliov
2026-04-21 18:07:17 +03:00
parent 6ad9c4d149
commit 998d52cb57
4 changed files with 210 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
src/datasets/gtauav_dataset.py
View File

@@ -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:
return self.image_transform(rgb)
if transform is not None:
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:

64
src/losses/multi_infonce.py
View File

@@ -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_final: Temperature after cosine decay.
temperature_init: Initial temperature value.
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,12 +117,19 @@ class InfoNCELoss(nn.Module):
Returns:
Dict with 'total', 'temperature', 'gate'.
"""
tau = cosine_temperature(
epoch=epoch,
total_epochs=total_epochs,
tau_init=self.temperature_init,
tau_final=self.temperature_final,
)
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,
tau_init=self.temperature_init,
tau_final=self.temperature_final,
)
loss = _symmetric_info_nce(
emb_a=embeddings["query"],
@@ -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),
}

53
src/models/asymmetric_encoder.py
View File

@@ -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(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
])
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),
])

115
src/training/train_gtauav.py
View File

@@ -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_final: float = 0.01
tau_init: float = 0.07
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.
transform = get_dino_transform(image_size=256)
# Data — separate transforms for train (augmented) and eval (clean).
drone_train_tf = get_drone_train_transform(image_size=256)
sat_train_tf = get_satellite_train_transform(image_size=256)
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 = AdamW(
model.trainable_parameters(),
lr=cfg.learning_rate,
weight_decay=cfg.weight_decay,
# Optimizer — per-group LR (text encoder gets lower LR).
param_groups = _build_param_groups(model, cfg.learning_rate, cfg.text_lr_factor)
# Include loss temperature if learnable.
if cfg.learnable_temperature and loss_fn.logit_scale is not None:
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: