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

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: