diff --git a/conf/gtauav_balanced_stripnet.gin b/conf/gtauav_balanced_stripnet.gin new file mode 100644 index 0000000..522776a --- /dev/null +++ b/conf/gtauav_balanced_stripnet.gin @@ -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" diff --git a/conf/gtauav_baseline_stripnet.gin b/conf/gtauav_baseline_stripnet.gin new file mode 100644 index 0000000..1974fa1 --- /dev/null +++ b/conf/gtauav_baseline_stripnet.gin @@ -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 diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 275ed37..872165e 100644 --- a/src/models/asymmetric_encoder.py +++ b/src/models/asymmetric_encoder.py @@ -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.dgtrs.model import DGTRSTextEncoder, load_dgtrs_text_encoder, tokenize_dgtrs 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, lora_rank: int = 4, device: str = "cuda", + backbone: str = "dinov3", + stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth", + stripnet_mona_last_n_stages: int = 2, ) -> None: 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.shared_encoder = shared_encoder + self.backbone = backbone self.device = device # 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._freeze(self.image_encoder) inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks) diff --git a/src/models/stripnet/__init__.py b/src/models/stripnet/__init__.py new file mode 100644 index 0000000..c2648be --- /dev/null +++ b/src/models/stripnet/__init__.py @@ -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", +] diff --git a/src/models/stripnet/conv_mona.py b/src/models/stripnet/conv_mona.py new file mode 100644 index 0000000..c13c0eb --- /dev/null +++ b/src/models/stripnet/conv_mona.py @@ -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 diff --git a/src/models/stripnet/model.py b/src/models/stripnet/model.py new file mode 100644 index 0000000..ff8c108 --- /dev/null +++ b/src/models/stripnet/model.py @@ -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 diff --git a/src/models/stripnet_encoder.py b/src/models/stripnet_encoder.py new file mode 100644 index 0000000..c0a2dc9 --- /dev/null +++ b/src/models/stripnet_encoder.py @@ -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] diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 6590f68..7c51588 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -90,6 +90,10 @@ class TrainConfigGTAUAV: mona_bottleneck: int = 64 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" + 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. resume_from: str | None = None # path to checkpoint for resuming @@ -564,7 +568,10 @@ def train(cfg: TrainConfigGTAUAV) -> None: start_epoch = resume_ckpt.get("epoch", -1) + 1 else: mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)" - enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)" + 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)" LOGGER.info("Building model — %s, %s", mode_str, enc_str) model = AsymmetricEncoder( dino_web_path=cfg.dino_web_path, @@ -576,11 +583,15 @@ def train(cfg: TrainConfigGTAUAV) -> None: 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, ).to(cfg.device) LOGGER.info("embed_dim=%d", model.embed_dim) # --- 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: model.image_encoder.set_gradient_checkpointing(True) else: @@ -589,6 +600,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": + 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_total = sum(p.numel() for p in model.parameters())