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

add sofia models

Browse Source
This commit is contained in:
pikaliov
2026-04-29 08:04:33 +03:00
parent a27b5a7357
commit 0d8c82acc3
23 changed files with 4448 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

188
src/training/train_gtauav.py
View File

@@ -54,6 +54,14 @@ from src.models.asymmetric_encoder import (
get_drone_train_transform,
get_satellite_train_transform,
)
from src.models.sofia_fusion_encoder import SOFIAFusionEncoder
from src.models.sofia_v1 import SOFIAv1Config
from src.models.sofia_v1_fusion_encoder import SOFIAv1FusionEncoder
from src.models.sofia_v71 import (
sofia_l_config,
sofia_m_config,
sofia_tiny_config,
)
LOGGER = logging.getLogger("caption_test.train_gtauav")
@@ -91,11 +99,30 @@ class TrainConfigGTAUAV:
mona_last_n_blocks: int = 12 # inject adapters only in last 12 of 24 ViT blocks
gradient_checkpointing: bool = True # trade compute for VRAM (allows larger batch)
# StripNet backbone option (replaces DINOv3 when backbone="stripnet").
backbone: str = "dinov3" # "dinov3" or "stripnet"
backbone: str = "dinov3" # "dinov3", "stripnet", or "sofia"
stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth"
stripnet_mona_last_n_stages: int = 2 # Conv-MONA in last N of 4 StripNet stages (0 = disable MONA)
stripnet_freeze: bool = True # If False, StripNet backbone is fully trainable (full fine-tune)
stripnet_backbone_lr_factor: float = 0.1 # Backbone LR = learning_rate * factor (only when unfrozen)
# SOFIA backbone options (used when backbone="sofia"). Trained from scratch — no pretrained checkpoint.
sofia_preset: str = "Tiny" # "Tiny" | "M" | "L"
sofia_d_descriptor: int = 1024 # retrieval space (1024 = match TextFusionMLP out_dim)
sofia_use_text_film_uav: bool = True # mid-level text-FiLM in UAV head
sofia_use_text_film_sat: bool = True # mid-level text-FiLM in SAT head
sofia_lora_rank: int = 4
sofia_mamba_variant: str = "mamba2" # "mamba1" | "mamba2" | "efficient_vmamba"
sofia_mamba_backend: str = "auto" # "auto" | "torch" | "mamba_ssm"
# EVSSBridge (B6-inspired refinement between heterogeneous stages, opt-in).
sofia_use_evss_bridge: bool = False
sofia_evss_bridge_locations: list[str] = field(default_factory=lambda: ["pre_stage3"])
# SOFIA v1 backbone options (used when backbone="sofia_v1"). StripNet+DCN, from scratch.
sofia_v1_variant: str = "small" # "tiny_tiny" | "tiny" | "small" | "small_v2"
sofia_v1_dcn_variant: str = "v2" # "v2" (torchvision DeformConv2d, stable) | "v4" (OpenGVLab, leaky)
sofia_v1_d_descriptor: int = 1024
sofia_v1_use_text_film_uav: bool = True
sofia_v1_use_text_film_sat: bool = True
sofia_v1_use_film_altitude: bool = True
sofia_v1_lora_rank: int = 4
# Training.
resume_from: str | None = None # path to checkpoint for resuming
@@ -167,7 +194,7 @@ def _atomic_save(obj: dict, path: Path) -> None:
def _build_param_groups(
model: AsymmetricEncoder,
model: nn.Module,
lr: float,
text_lr_factor: float,
stripnet_backbone_lr_factor: float = 0.1,
@@ -177,7 +204,8 @@ def _build_param_groups(
Groups:
- text_encoder.* → lr * text_lr_factor (default 1e-5)
- image_encoder.backbone.* (when StripNet unfrozen) → lr * stripnet_backbone_lr_factor (default 1e-5)
- everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv) → lr
- everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv,
SOFIA backbone+heads when backbone="sofia") → lr
"""
text_params = []
backbone_params = []
@@ -249,9 +277,13 @@ def _embed_drone_queries(
embs: list[torch.Tensor] = []
for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False):
drone_img = batch["drone_img"].to(device, non_blocking=True)
altitude = batch.get("altitude")
if altitude is not None:
altitude = altitude.to(device, non_blocking=True)
q = model.encode_query(
drone_img,
batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
altitude=altitude,
)
embs.append(q.cpu())
@@ -336,9 +368,13 @@ def _evaluate(
if max_batches is not None and i >= max_batches:
break
drone_img = batch["drone_img"].to(device, non_blocking=True)
altitude = batch.get("altitude")
if altitude is not None:
altitude = altitude.to(device, non_blocking=True)
q = model.encode_query(
drone_img,
batch["caption_l1"], batch["caption_l2"], batch["caption_l3"],
altitude=altitude,
)
query_embs.append(q.cpu())
query_valid_names.extend(batch["valid_sat_names"])
@@ -575,40 +611,97 @@ def train(cfg: TrainConfigGTAUAV) -> None:
if cfg.resume_from is not None:
LOGGER.info("Resuming from %s", cfg.resume_from)
model, resume_ckpt = AsymmetricEncoder.load_checkpoint(
cfg.resume_from,
dino_web_path=cfg.dino_web_path,
dino_sat_path=cfg.dino_sat_path,
lrsclip_path=cfg.lrsclip_path,
device=cfg.device,
)
if cfg.backbone == "sofia":
model, resume_ckpt = SOFIAFusionEncoder.load_checkpoint(
cfg.resume_from,
lrsclip_path=cfg.lrsclip_path,
device=cfg.device,
)
elif cfg.backbone == "sofia_v1":
model, resume_ckpt = SOFIAv1FusionEncoder.load_checkpoint(
cfg.resume_from,
lrsclip_path=cfg.lrsclip_path,
device=cfg.device,
)
else:
model, resume_ckpt = AsymmetricEncoder.load_checkpoint(
cfg.resume_from,
dino_web_path=cfg.dino_web_path,
dino_sat_path=cfg.dino_sat_path,
lrsclip_path=cfg.lrsclip_path,
device=cfg.device,
)
start_epoch = resume_ckpt.get("epoch", -1) + 1
else:
mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)"
if cfg.backbone == "stripnet":
if cfg.backbone == "sofia":
enc_str = f"SOFIA-{cfg.sofia_preset} (text-FiLM uav={cfg.sofia_use_text_film_uav}, sat={cfg.sofia_use_text_film_sat})"
elif cfg.backbone == "sofia_v1":
enc_str = f"SOFIAv1-{cfg.sofia_v1_variant} (StripNet+DCNv4, text-FiLM uav={cfg.sofia_v1_use_text_film_uav}, sat={cfg.sofia_v1_use_text_film_sat})"
elif cfg.backbone == "stripnet":
enc_str = "StripNet-small (shared, 512→1024 proj)"
else:
enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)"
LOGGER.info("Building model — %s, %s", mode_str, enc_str)
model = AsymmetricEncoder(
dino_web_path=cfg.dino_web_path,
dino_sat_path=cfg.dino_sat_path,
lrsclip_path=cfg.lrsclip_path,
init_gate=cfg.init_gate,
baseline_mode=cfg.baseline_mode,
shared_encoder=cfg.shared_encoder,
mona_bottleneck=cfg.mona_bottleneck,
mona_last_n_blocks=cfg.mona_last_n_blocks,
device=cfg.device,
backbone=cfg.backbone,
stripnet_path=cfg.stripnet_path,
stripnet_mona_last_n_stages=cfg.stripnet_mona_last_n_stages,
stripnet_freeze=cfg.stripnet_freeze,
).to(cfg.device)
if cfg.backbone == "sofia":
preset_map = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config}
if cfg.sofia_preset not in preset_map:
raise ValueError(f"Unknown sofia_preset={cfg.sofia_preset!r}")
sofia_cfg = preset_map[cfg.sofia_preset]()
sofia_cfg.d_descriptor = cfg.sofia_d_descriptor
sofia_cfg.text_film_dim = cfg.sofia_d_descriptor
sofia_cfg.use_text_film_uav = cfg.sofia_use_text_film_uav and not cfg.baseline_mode
sofia_cfg.use_text_film_sat = cfg.sofia_use_text_film_sat and not cfg.baseline_mode
sofia_cfg.mamba_variant = cfg.sofia_mamba_variant
sofia_cfg.mamba_backend = cfg.sofia_mamba_backend
sofia_cfg.use_evss_bridge = cfg.sofia_use_evss_bridge
sofia_cfg.evss_bridge_locations = list(cfg.sofia_evss_bridge_locations)
model = SOFIAFusionEncoder(
sofia_cfg=sofia_cfg,
lrsclip_path=cfg.lrsclip_path,
init_gate=cfg.init_gate,
baseline_mode=cfg.baseline_mode,
lora_rank=cfg.sofia_lora_rank,
device=cfg.device,
).to(cfg.device)
elif cfg.backbone == "sofia_v1":
sofia_v1_cfg = SOFIAv1Config(
variant=cfg.sofia_v1_variant,
dcn_variant=cfg.sofia_v1_dcn_variant,
d_descriptor=cfg.sofia_v1_d_descriptor,
text_film_dim=cfg.sofia_v1_d_descriptor,
use_text_film_uav=cfg.sofia_v1_use_text_film_uav and not cfg.baseline_mode,
use_text_film_sat=cfg.sofia_v1_use_text_film_sat and not cfg.baseline_mode,
use_film_altitude=cfg.sofia_v1_use_film_altitude,
)
model = SOFIAv1FusionEncoder(
sofia_cfg=sofia_v1_cfg,
lrsclip_path=cfg.lrsclip_path,
init_gate=cfg.init_gate,
baseline_mode=cfg.baseline_mode,
lora_rank=cfg.sofia_v1_lora_rank,
device=cfg.device,
).to(cfg.device)
else:
model = AsymmetricEncoder(
dino_web_path=cfg.dino_web_path,
dino_sat_path=cfg.dino_sat_path,
lrsclip_path=cfg.lrsclip_path,
init_gate=cfg.init_gate,
baseline_mode=cfg.baseline_mode,
shared_encoder=cfg.shared_encoder,
mona_bottleneck=cfg.mona_bottleneck,
mona_last_n_blocks=cfg.mona_last_n_blocks,
device=cfg.device,
backbone=cfg.backbone,
stripnet_path=cfg.stripnet_path,
stripnet_mona_last_n_stages=cfg.stripnet_mona_last_n_stages,
stripnet_freeze=cfg.stripnet_freeze,
).to(cfg.device)
LOGGER.info("embed_dim=%d", model.embed_dim)
# --- Gradient checkpointing (trade compute for VRAM) ---
# StripNet doesn't expose set_gradient_checkpointing — skip silently.
# StripNet/SOFIA don't expose set_gradient_checkpointing — only DGTRS gets it.
if cfg.gradient_checkpointing and cfg.backbone == "dinov3":
if cfg.shared_encoder:
model.image_encoder.set_gradient_checkpointing(True)
@@ -618,10 +711,10 @@ def train(cfg: TrainConfigGTAUAV) -> None:
if model.text_encoder is not None:
model.text_encoder.transformer.gradient_checkpointing = True
LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)")
elif cfg.gradient_checkpointing and cfg.backbone == "stripnet":
elif cfg.gradient_checkpointing and cfg.backbone in ("stripnet", "sofia", "sofia_v1"):
if model.text_encoder is not None:
model.text_encoder.transformer.gradient_checkpointing = True
LOGGER.info("Gradient checkpointing enabled (DGTRS only; StripNet doesn't support)")
LOGGER.info("Gradient checkpointing enabled (DGTRS only; %s doesn't support)", cfg.backbone)
n_trainable = sum(p.numel() for p in model.trainable_parameters())
n_total = sum(p.numel() for p in model.parameters())
@@ -879,11 +972,14 @@ def train(cfg: TrainConfigGTAUAV) -> None:
drone_img = batch["drone_img"].to(cfg.device, non_blocking=True)
sat_img = batch["sat_img"].to(cfg.device, non_blocking=True)
altitude = batch.get("altitude")
if altitude is not None:
altitude = altitude.to(cfg.device, non_blocking=True)
# Model forward in AMP (fp16 for DINOv3/DGTRS encoders).
with autocast(device_type="cuda", enabled=cfg.use_amp):
if cfg.baseline_mode:
embeddings = model(drone_img=drone_img, sat_img=sat_img)
embeddings = model(drone_img=drone_img, sat_img=sat_img, altitude=altitude)
else:
embeddings = model(
drone_img=drone_img,
@@ -894,6 +990,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
sat_caption_l1=batch["sat_caption_l1"],
sat_caption_l2=batch["sat_caption_l2"],
sat_caption_l3=batch["sat_caption_l3"],
altitude=altitude,
)
# Loss — InfoNCE or WeightedInfoNCE. Only the latter uses positive_weights.
queue_neg = neg_bank.get_queue() if neg_bank is not None else None
@@ -1103,20 +1200,23 @@ def train(cfg: TrainConfigGTAUAV) -> None:
history.append(epoch_record)
# Save checkpoint. Model architecture flags go into the ckpt so
# `AsymmetricEncoder.load_checkpoint` can rebuild the right shape.
_atomic_save(
obj={
"epoch": epoch,
"model_state": model.state_dict(),
"optimizer_state": optimizer.state_dict(),
"loss_state": loss_fn.state_dict(),
"baseline_mode": cfg.baseline_mode,
"shared_encoder": cfg.shared_encoder,
"mona_bottleneck": cfg.mona_bottleneck,
"mona_last_n_blocks": cfg.mona_last_n_blocks,
},
path=output_dir / f"ckpt_epoch{epoch:03d}.pt",
)
# `AsymmetricEncoder.load_checkpoint` (or `SOFIAFusionEncoder.load_checkpoint`)
# can rebuild the right shape.
ckpt_obj = {
"epoch": epoch,
"model_state": model.state_dict(),
"optimizer_state": optimizer.state_dict(),
"loss_state": loss_fn.state_dict(),
"baseline_mode": cfg.baseline_mode,
"backbone": cfg.backbone,
}
if cfg.backbone in ("sofia", "sofia_v1"):
ckpt_obj["sofia_cfg"] = model.sofia_cfg
else:
ckpt_obj["shared_encoder"] = cfg.shared_encoder
ckpt_obj["mona_bottleneck"] = cfg.mona_bottleneck
ckpt_obj["mona_last_n_blocks"] = cfg.mona_last_n_blocks
_atomic_save(obj=ckpt_obj, path=output_dir / f"ckpt_epoch{epoch:03d}.pt")
LOGGER.info("Checkpoint saved: ckpt_epoch%03d.pt", epoch)
# Save history.