Add StripNet backbone option with Conv-MONA adaptation

StripNet-small (Strip-R-CNN, HVision-NKU) as alternative image encoder to
DINOv3 ViT-L/16. ~28M params (10x smaller). Output 512-dim from stage 4
projected to 1024 to keep retrieval space compatible with DINOv3 configs.

- src/models/stripnet/: self-contained backbone (model.py, conv_mona.py).
  State-dict naming follows upstream Strip-R-CNN repo (conv_spatial1/2);
  ImageNet-1K pretrained head dropped on load.
- Conv-MONA: 2D adaptation of MONA (CVPR 2025) for CNN blocks. BN → 1x1
  Down(C->bn) → multi-scale DWConv {3,5,7} mean → +residual → GELU →
  1x1 Up(bn->C) with channel-wise layer scale γ init 1e-6. Two adapters
  per StripNet Block (post-attn, post-mlp); injected into deepest N stages.
- StripNetEncoder: GAP + Linear(512->1024). Overrides train() to keep
  frozen BatchNorm stats stable across mode flips.
- AsymmetricEncoder: new `backbone="stripnet"` option (always shared).
- TrainConfigGTAUAV: backbone, stripnet_path, stripnet_mona_last_n_stages.
- conf/gtauav_balanced_stripnet.gin + gtauav_baseline_stripnet.gin.

Smoke test: forward [2,3,256,256] -> [2,1024]. Trainable: 1.2M baseline
(8.27%), 4.76M with text (3.35% of 142M).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-25 14:34:53 +03:00
parent 814586ce3b
commit d4cb2dd300
8 changed files with 487 additions and 4 deletions

View File

@@ -0,0 +1,27 @@
# GTA-UAV Balanced (StripNet backbone): StripNet-small + Conv-MONA in last 2 stages.
# Replaces DINOv3 ViT-L/16 with strip-shaped DWConv hierarchical CNN (~28M params,
# 10× smaller than DINOv3). Output 512-dim → projected to 1024 to match retrieval space.
#
# Trainable:
# - Projection (Linear 512→1024): ~525K
# - Conv-MONA in stages 3 & 4 (2 adapters per Block × 6 blocks total): ~2-3M
# - LoRA on DGTRS-CLIP: 147K
# - TextFusionMLP (shared): ~3.4M
# - GatedFusion gates + tau: 3 scalars
#
# StripNet pretrained on ImageNet-1K (head dropped); state-dict naming follows
# upstream Strip-R-CNN repo (`conv_spatial1/2`).
include 'conf/gtauav_balanced.gin'
# ---- Backbone ----
TrainConfigGTAUAV.backbone = "stripnet"
TrainConfigGTAUAV.stripnet_path = "nn_models/STRIPNET/stripnet_s.pth"
TrainConfigGTAUAV.stripnet_mona_last_n_stages = 2 # Conv-MONA in stages 3 & 4 (deepest)
# ---- Model overrides ----
TrainConfigGTAUAV.shared_encoder = True # StripNet always shared (one CNN for both branches)
TrainConfigGTAUAV.mona_bottleneck = 64 # Conv-MONA bottleneck channels
# ---- Output ----
TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_stripnet"

View File

@@ -0,0 +1,8 @@
# GTA-UAV Baseline (StripNet backbone): no text fusion. Reference R@1 for
# computing Δ R@1 against gtauav_balanced_stripnet.gin.
include 'conf/gtauav_balanced_stripnet.gin'
TrainConfigGTAUAV.baseline_mode = True
TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_stripnet"
TrainConfigGTAUAV.use_gradcam = False

View File

@@ -28,6 +28,8 @@ from safetensors.torch import load_file as load_safetensors
from src.models.adapters import inject_lora_into_dgtrs, inject_mona_into_dinov3 from src.models.adapters import inject_lora_into_dgtrs, inject_mona_into_dinov3
from src.models.dgtrs.model import DGTRSTextEncoder, load_dgtrs_text_encoder, tokenize_dgtrs from src.models.dgtrs.model import DGTRSTextEncoder, load_dgtrs_text_encoder, tokenize_dgtrs
from src.models.dual_encoder import GatedFusion, ProjectionHead from src.models.dual_encoder import GatedFusion, ProjectionHead
from src.models.stripnet import inject_conv_mona_into_stripnet
from src.models.stripnet_encoder import StripNetEncoder
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -302,15 +304,30 @@ class AsymmetricEncoder(nn.Module):
mona_last_n_blocks: int = 24, mona_last_n_blocks: int = 24,
lora_rank: int = 4, lora_rank: int = 4,
device: str = "cuda", device: str = "cuda",
backbone: str = "dinov3",
stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth",
stripnet_mona_last_n_stages: int = 2,
) -> None: ) -> None:
super().__init__() super().__init__()
self.embed_dim = self.DINO_DIM # native 1024, no projection self.embed_dim = self.DINO_DIM # native 1024 (StripNet projects 512 -> 1024)
self.baseline_mode = baseline_mode self.baseline_mode = baseline_mode
self.shared_encoder = shared_encoder self.shared_encoder = shared_encoder
self.backbone = backbone
self.device = device self.device = device
# Image encoder(s) (frozen + MONA adapters). # Image encoder(s) (frozen + MONA adapters).
if shared_encoder: if backbone == "stripnet":
# StripNet always operates as shared encoder (one CNN for both branches).
self.shared_encoder = True
self.image_encoder = StripNetEncoder(checkpoint_path=stripnet_path, out_dim=self.DINO_DIM)
self._freeze(self.image_encoder.backbone)
inject_conv_mona_into_stripnet(
self.image_encoder.backbone,
bottleneck=mona_bottleneck,
last_n_stages=stripnet_mona_last_n_stages,
)
LOGGER.info("StripNet backbone: shared encoder, projection 512 -> %d", self.DINO_DIM)
elif shared_encoder:
self.image_encoder = DINOv3ViT.from_pretrained(dino_web_path) self.image_encoder = DINOv3ViT.from_pretrained(dino_web_path)
self._freeze(self.image_encoder) self._freeze(self.image_encoder)
inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks) inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)

View File

@@ -0,0 +1,9 @@
from src.models.stripnet.model import StripNet, get_stripnet_small, load_stripnet_small_pretrained
from src.models.stripnet.conv_mona import inject_conv_mona_into_stripnet
__all__ = [
"StripNet",
"get_stripnet_small",
"load_stripnet_small_pretrained",
"inject_conv_mona_into_stripnet",
]

View File

@@ -0,0 +1,106 @@
"""Conv-MONA: 2D adaptation of MONA (CVPR 2025) for hierarchical CNN backbones.
MONA paper applies sequence-form adapters after MSA / MLP in ViT blocks. Here we
mirror that idea in [B, C, H, W] form: BN → 1×1 Down(C→bn) → multi-scale DWConv
{3,5,7} mean → +residual → GELU → 1×1 Up(bn→C). Layer scale (γ) channel-wise,
init 1e-6 for near-identity start. Two adapters per StripNet Block: post-attn
and post-mlp.
"""
from __future__ import annotations
import logging
import torch
import torch.nn as nn
from src.models.stripnet.model import StripNet, Block
LOGGER = logging.getLogger("caption_test.stripnet.adapters")
class ConvMona(nn.Module):
"""Single Conv-MONA adapter.
Args:
dim: input channel dim.
bottleneck: bottleneck channel dim (e.g. 64).
gamma_init: layer-scale init value (1e-6 for near-identity at start).
"""
def __init__(self, dim: int, bottleneck: int = 64, gamma_init: float = 1e-6) -> None:
super().__init__()
self.norm = nn.BatchNorm2d(dim)
self.down = nn.Conv2d(dim, bottleneck, kernel_size=1, bias=True)
self.dw3 = nn.Conv2d(bottleneck, bottleneck, kernel_size=3, padding=1, groups=bottleneck, bias=True)
self.dw5 = nn.Conv2d(bottleneck, bottleneck, kernel_size=5, padding=2, groups=bottleneck, bias=True)
self.dw7 = nn.Conv2d(bottleneck, bottleneck, kernel_size=7, padding=3, groups=bottleneck, bias=True)
self.act = nn.GELU()
self.up = nn.Conv2d(bottleneck, dim, kernel_size=1, bias=True)
# Channel-wise layer scale (γ), broadcast across H, W.
self.gamma = nn.Parameter(gamma_init * torch.ones(dim), requires_grad=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
h = self.norm(x)
h = self.down(h)
h = (self.dw3(h) + self.dw5(h) + self.dw7(h)) / 3.0 + h
h = self.act(h)
h = self.up(h)
return self.gamma.view(1, -1, 1, 1) * h
def _patched_block_forward(block: Block, mona_attn: ConvMona, mona_mlp: ConvMona):
"""Closure that wraps a Block.forward with two Conv-MONA residuals."""
orig_attn = block.attn
orig_mlp = block.mlp
orig_norm1 = block.norm1
orig_norm2 = block.norm2
orig_drop = block.drop_path
ls1 = block.layer_scale_1
ls2 = block.layer_scale_2
def forward(x: torch.Tensor) -> torch.Tensor:
x = x + orig_drop(ls1.unsqueeze(-1).unsqueeze(-1) * orig_attn(orig_norm1(x))) + mona_attn(x)
x = x + orig_drop(ls2.unsqueeze(-1).unsqueeze(-1) * orig_mlp(orig_norm2(x))) + mona_mlp(x)
return x
return forward
def inject_conv_mona_into_stripnet(
model: StripNet,
bottleneck: int = 64,
last_n_stages: int = 2,
use_bf16: bool = False,
) -> int:
"""Inject Conv-MONA adapters into the deepest `last_n_stages` of StripNet.
Each Block in the targeted stages gets two adapters (post-attn, post-mlp).
Returns the number of adapters injected.
Stages are 1-indexed in StripNet (block1..block4). With `last_n_stages=2`
we adapt block3 and block4 — the deepest, semantically richest features.
"""
n_stages = model.num_stages
target_stages = list(range(max(1, n_stages - last_n_stages + 1), n_stages + 1))
n_added = 0
for stage_idx in target_stages:
blocks: nn.ModuleList = getattr(model, f"block{stage_idx}")
dim = model.embed_dims[stage_idx - 1]
for blk_idx, block in enumerate(blocks):
mona_a = ConvMona(dim=dim, bottleneck=bottleneck)
mona_m = ConvMona(dim=dim, bottleneck=bottleneck)
if use_bf16:
mona_a.to(dtype=torch.bfloat16)
mona_m.to(dtype=torch.bfloat16)
# Register as submodules so they get moved by .to(device) / .train() etc.
block.add_module(f"mona_attn", mona_a)
block.add_module(f"mona_mlp", mona_m)
block.forward = _patched_block_forward(block, mona_a, mona_m)
n_added += 2
n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
LOGGER.info(
"🔧 Conv-MONA injected: %d adapters in stages %s, %d trainable params (bottleneck=%d)",
n_added, target_stages, n_trainable, bottleneck,
)
return n_added

View File

@@ -0,0 +1,253 @@
"""StripNet (small) backbone — adapted from Strip-R-CNN (HVision-NKU).
Self-contained: no external utils. State-dict naming follows the upstream
ImageNet-pretrained checkpoint (`conv_spatial1/2` for the strip kernels).
"""
from __future__ import annotations
import logging
import math
from functools import partial
from pathlib import Path
from typing import List
import torch
import torch.nn as nn
import torch.nn.functional as F
LOGGER = logging.getLogger("caption_test.stripnet")
def _to_2tuple(x):
if isinstance(x, (tuple, list)):
return tuple(x)
return (x, x)
def drop_path(x: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
if drop_prob == 0.0 or not training:
return x
keep = 1.0 - drop_prob
shape = (x.shape[0],) + (1,) * (x.ndim - 1)
rand = x.new_empty(shape).bernoulli_(keep)
if keep > 0:
rand.div_(keep)
return x * rand
class DropPath(nn.Module):
def __init__(self, p: float = 0.0) -> None:
super().__init__()
self.p = p
def forward(self, x: torch.Tensor) -> torch.Tensor:
return drop_path(x, self.p, self.training)
class DWConv(nn.Module):
def __init__(self, dim: int) -> None:
super().__init__()
self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.dwconv(x)
class Mlp(nn.Module):
def __init__(self, in_features: int, hidden_features: int, drop: float = 0.0) -> None:
super().__init__()
self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
self.dwconv = DWConv(hidden_features)
self.act = nn.GELU()
self.fc2 = nn.Conv2d(hidden_features, in_features, 1)
self.drop = nn.Dropout(drop)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.fc1(x)
x = self.dwconv(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class StripGatingUnit(nn.Module):
"""Strip spatial gating: 5x5 DWConv -> (1, k2) -> (k2, 1) -> 1x1 -> gate."""
def __init__(self, dim: int, k1: int, k2: int) -> None:
super().__init__()
self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
# Names match upstream pretrained checkpoint: conv_spatial1 / conv_spatial2.
self.conv_spatial1 = nn.Conv2d(dim, dim, kernel_size=(k1, k2), stride=1,
padding=(k1 // 2, k2 // 2), groups=dim)
self.conv_spatial2 = nn.Conv2d(dim, dim, kernel_size=(k2, k1), stride=1,
padding=(k2 // 2, k1 // 2), groups=dim)
self.conv1 = nn.Conv2d(dim, dim, 1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
attn = self.conv0(x)
attn = self.conv_spatial1(attn)
attn = self.conv_spatial2(attn)
attn = self.conv1(attn)
return x * attn
class StripAttention(nn.Module):
def __init__(self, dim: int, k1: int, k2: int) -> None:
super().__init__()
self.proj_1 = nn.Conv2d(dim, dim, 1)
self.activation = nn.GELU()
self.spatial_gating_unit = StripGatingUnit(dim, k1, k2)
self.proj_2 = nn.Conv2d(dim, dim, 1)
def forward(self, x: torch.Tensor) -> torch.Tensor:
residual = x
x = self.proj_1(x)
x = self.activation(x)
x = self.spatial_gating_unit(x)
x = self.proj_2(x)
return x + residual
class Block(nn.Module):
def __init__(self, dim: int, mlp_ratio: float, k1: int, k2: int, drop: float, drop_path: float) -> None:
super().__init__()
self.norm1 = nn.BatchNorm2d(dim)
self.norm2 = nn.BatchNorm2d(dim)
self.attn = StripAttention(dim, k1, k2)
self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
self.mlp = Mlp(dim, int(dim * mlp_ratio), drop=drop)
ls_init = 1e-2
self.layer_scale_1 = nn.Parameter(ls_init * torch.ones(dim), requires_grad=True)
self.layer_scale_2 = nn.Parameter(ls_init * torch.ones(dim), requires_grad=True)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x + self.drop_path(
self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * self.attn(self.norm1(x))
)
x = x + self.drop_path(
self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * self.mlp(self.norm2(x))
)
return x
class OverlapPatchEmbed(nn.Module):
def __init__(self, patch_size: int, stride: int, in_chans: int, embed_dim: int) -> None:
super().__init__()
ph, pw = _to_2tuple(patch_size)
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=(ph, pw), stride=stride,
padding=(ph // 2, pw // 2))
self.norm = nn.BatchNorm2d(embed_dim)
def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int, int]:
x = self.proj(x)
_, _, H, W = x.shape
x = self.norm(x)
return x, H, W
class StripNet(nn.Module):
"""Strip-R-CNN backbone: 4-stage hierarchical CNN with strip-shaped DWConv attention.
Output: list of [B, C_i, H/s_i, W/s_i] per stage. Use `forward_last_features` for
the deepest stage only.
"""
def __init__(
self,
embed_dims: List[int] = [64, 128, 320, 512],
mlp_ratios: List[int] = [8, 8, 4, 4],
k1s: List[int] = [1, 1, 1, 1],
k2s: List[int] = [19, 19, 19, 19],
depths: List[int] = [2, 2, 4, 2],
drop_rate: float = 0.1,
drop_path_rate: float = 0.15,
in_chans: int = 3,
) -> None:
super().__init__()
self.depths = depths
self.num_stages = len(depths)
self.embed_dims = embed_dims
dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
cur = 0
for i in range(self.num_stages):
patch_embed = OverlapPatchEmbed(
patch_size=7 if i == 0 else 3,
stride=4 if i == 0 else 2,
in_chans=in_chans if i == 0 else embed_dims[i - 1],
embed_dim=embed_dims[i],
)
block = nn.ModuleList([
Block(dim=embed_dims[i], mlp_ratio=mlp_ratios[i], k1=k1s[i], k2=k2s[i],
drop=drop_rate, drop_path=dpr[cur + j])
for j in range(depths[i])
])
norm = nn.LayerNorm(embed_dims[i], eps=1e-6)
cur += depths[i]
setattr(self, f"patch_embed{i + 1}", patch_embed)
setattr(self, f"block{i + 1}", block)
setattr(self, f"norm{i + 1}", norm)
def forward_features(self, x: torch.Tensor) -> List[torch.Tensor]:
B = x.shape[0]
outs: List[torch.Tensor] = []
for i in range(self.num_stages):
patch_embed = getattr(self, f"patch_embed{i + 1}")
block = getattr(self, f"block{i + 1}")
norm = getattr(self, f"norm{i + 1}")
x, H, W = patch_embed(x)
for blk in block:
x = blk(x)
x = x.flatten(2).transpose(1, 2)
x = norm(x)
x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
outs.append(x)
return outs
def forward_last_features(self, x: torch.Tensor) -> torch.Tensor:
return self.forward_features(x)[-1]
def get_stripnet_small() -> StripNet:
return StripNet(
embed_dims=[64, 128, 320, 512],
mlp_ratios=[8, 8, 4, 4],
k1s=[1, 1, 1, 1],
k2s=[19, 19, 19, 19],
depths=[2, 2, 4, 2],
drop_rate=0.1,
drop_path_rate=0.15,
)
def load_stripnet_small_pretrained(checkpoint_path: str | Path) -> StripNet:
"""Build StripNet-small and load ImageNet-pretrained weights.
Strips the classification `head.*` keys. Tolerates missing/extra keys
(norm{N}.* are LayerNorm here vs BatchNorm in some forks — we keep LN).
"""
LOGGER.info("📐 Loading StripNet-small from %s", checkpoint_path)
model = get_stripnet_small()
raw = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
state = raw.get("state_dict", raw) if isinstance(raw, dict) else raw
# Drop classification head + the BN-form norm{N} keys if present (we use LN here).
drop_prefixes = ("head.",)
cleaned = {k: v for k, v in state.items() if not any(k.startswith(p) for p in drop_prefixes)}
# The pretrained checkpoint stores norm{N} as BatchNorm2d (running_mean/var/num_batches_tracked).
# Our code uses LayerNorm at this position. Strip BN running stats if found; copy weight/bias.
for n in (1, 2, 3, 4):
for suffix in ("running_mean", "running_var", "num_batches_tracked"):
cleaned.pop(f"norm{n}.{suffix}", None)
missing, unexpected = model.load_state_dict(cleaned, strict=False)
if missing:
LOGGER.info("StripNet missing keys (expected for newly-init layers): %d", len(missing))
if unexpected:
LOGGER.info("StripNet unexpected keys (ignored): %d", len(unexpected))
LOGGER.info("📐 StripNet-small loaded: %d params", sum(p.numel() for p in model.parameters()))
return model

View File

@@ -0,0 +1,48 @@
"""StripNet image encoder wrapper for the caption-test pipeline.
Exposes the same interface as DINOv3ViT: `forward(images) -> [B, embed_dim]`.
StripNet's deepest stage produces [B, 512, H/32, W/32]; we apply global average
pooling (GAP) and project to the target retrieval dimension via Linear(512→1024)
to match DINOv3 native dim and keep TextFusionMLP unchanged.
"""
from __future__ import annotations
import logging
import torch
import torch.nn as nn
from src.models.stripnet import StripNet, load_stripnet_small_pretrained
LOGGER = logging.getLogger("caption_test.stripnet_encoder")
class StripNetEncoder(nn.Module):
"""StripNet-small + GAP + projection to `out_dim`.
Frozen backbone (BatchNorm in eval mode); only the projection head and
any injected Conv-MONA adapters are trainable.
"""
LAST_STAGE_DIM = 512 # StripNet-small last stage embed dim
def __init__(self, checkpoint_path: str, out_dim: int = 1024) -> None:
super().__init__()
self.out_dim = out_dim
self.backbone: StripNet = load_stripnet_small_pretrained(checkpoint_path)
self.pool = nn.AdaptiveAvgPool2d(1)
self.projection = nn.Linear(self.LAST_STAGE_DIM, out_dim)
nn.init.trunc_normal_(self.projection.weight, std=0.02)
nn.init.zeros_(self.projection.bias)
def train(self, mode: bool = True):
"""Override: keep frozen backbone in eval mode (BN running stats stable)."""
super().train(mode)
# Frozen backbone always in eval; trainable adapters/projection follow `mode`.
if not any(p.requires_grad for p in self.backbone.parameters()):
self.backbone.eval()
return self
def forward(self, images: torch.Tensor) -> torch.Tensor:
feat = self.backbone.forward_last_features(images) # [B, 512, H/32, W/32]
pooled = self.pool(feat).flatten(1) # [B, 512]
return self.projection(pooled) # [B, out_dim]

View File

@@ -90,6 +90,10 @@ class TrainConfigGTAUAV:
mona_bottleneck: int = 64 mona_bottleneck: int = 64
mona_last_n_blocks: int = 12 # inject adapters only in last 12 of 24 ViT blocks 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) 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"
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
# Training. # Training.
resume_from: str | None = None # path to checkpoint for resuming resume_from: str | None = None # path to checkpoint for resuming
@@ -564,6 +568,9 @@ def train(cfg: TrainConfigGTAUAV) -> None:
start_epoch = resume_ckpt.get("epoch", -1) + 1 start_epoch = resume_ckpt.get("epoch", -1) + 1
else: else:
mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)" mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)"
if 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)" enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)"
LOGGER.info("Building model — %s, %s", mode_str, enc_str) LOGGER.info("Building model — %s, %s", mode_str, enc_str)
model = AsymmetricEncoder( model = AsymmetricEncoder(
@@ -576,11 +583,15 @@ def train(cfg: TrainConfigGTAUAV) -> None:
mona_bottleneck=cfg.mona_bottleneck, mona_bottleneck=cfg.mona_bottleneck,
mona_last_n_blocks=cfg.mona_last_n_blocks, mona_last_n_blocks=cfg.mona_last_n_blocks,
device=cfg.device, device=cfg.device,
backbone=cfg.backbone,
stripnet_path=cfg.stripnet_path,
stripnet_mona_last_n_stages=cfg.stripnet_mona_last_n_stages,
).to(cfg.device) ).to(cfg.device)
LOGGER.info("embed_dim=%d", model.embed_dim) LOGGER.info("embed_dim=%d", model.embed_dim)
# --- Gradient checkpointing (trade compute for VRAM) --- # --- Gradient checkpointing (trade compute for VRAM) ---
if cfg.gradient_checkpointing: # StripNet doesn't expose set_gradient_checkpointing — skip silently.
if cfg.gradient_checkpointing and cfg.backbone == "dinov3":
if cfg.shared_encoder: if cfg.shared_encoder:
model.image_encoder.set_gradient_checkpointing(True) model.image_encoder.set_gradient_checkpointing(True)
else: else:
@@ -589,6 +600,10 @@ def train(cfg: TrainConfigGTAUAV) -> None:
if model.text_encoder is not None: if model.text_encoder is not None:
model.text_encoder.transformer.gradient_checkpointing = True model.text_encoder.transformer.gradient_checkpointing = True
LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)") LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)")
elif cfg.gradient_checkpointing and cfg.backbone == "stripnet":
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)")
n_trainable = sum(p.numel() for p in model.trainable_parameters()) n_trainable = sum(p.numel() for p in model.trainable_parameters())
n_total = sum(p.numel() for p in model.parameters()) n_total = sum(p.numel() for p in model.parameters())