Add stripnet_freeze flag for full StripNet fine-tune mode

When stripnet_freeze=False, all StripNet backbone params train end-to-end
with a separate optimizer group at lr * stripnet_backbone_lr_factor (default
0.1, so 1e-5 with default learning_rate=1e-4) — typical fine-tuning practice
for ImageNet-pretrained CNNs to avoid catastrophic forgetting.

Conv-MONA is now optional (stripnet_mona_last_n_stages=0 disables it). Three
modes are now supported:
  - frozen + MONA: PEFT-style (~1.2M trainable, original default)
  - unfrozen, no MONA: full fine-tune (~13.85M, all backbone params)
  - unfrozen + MONA: hybrid (~14.5M, backbone + extra adapters)

_build_param_groups: new "backbone" group identifies image_encoder.backbone.*
params (excluding mona_*) when backbone="stripnet"; assigned lr factor
controls fine-tune step size independently from text/MONA groups.

conf/gtauav_balanced_stripnet_unfrozen.gin + baseline variant: ready-to-use
configs for full fine-tune experiment.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

conf/gtauav_balanced_stripnet_unfrozen.gin

@@ -0,0 +1,17 @@
+# GTA-UAV Balanced (StripNet, fully unfrozen): all StripNet layers trainable.
+# Backbone trains with reduced LR (lr * stripnet_backbone_lr_factor).
+# Conv-MONA disabled by default — full fine-tune supplies enough capacity.
+# Set stripnet_mona_last_n_stages > 0 if you want MONA + fine-tune hybrid.
+#
+# Note: StripNet uses BatchNorm. With small batch (8) and gradient accumulation,
+# BN running stats may drift. Watch validation loss for instability.
+
+include 'conf/gtauav_balanced_stripnet.gin'
+
+# ---- Unfreeze backbone ----
+TrainConfigGTAUAV.stripnet_freeze = False
+TrainConfigGTAUAV.stripnet_mona_last_n_stages = 0  # disable Conv-MONA (full fine-tune handles adaptation)
+TrainConfigGTAUAV.stripnet_backbone_lr_factor = 0.1  # backbone lr = 1e-4 * 0.1 = 1e-5
+
+# ---- Output ----
+TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_stripnet_unfrozen"

conf/gtauav_baseline_stripnet_unfrozen.gin

@@ -0,0 +1,8 @@
+# GTA-UAV Baseline (StripNet, fully unfrozen): no text fusion. For Δ R@1
+# against gtauav_balanced_stripnet_unfrozen.gin.
+
+include 'conf/gtauav_balanced_stripnet_unfrozen.gin'
+
+TrainConfigGTAUAV.baseline_mode = True
+TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_stripnet_unfrozen"
+TrainConfigGTAUAV.use_gradcam = False

src/models/asymmetric_encoder.py

@@ -307,6 +307,7 @@ class AsymmetricEncoder(nn.Module):
         backbone: str = "dinov3",
         stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth",
         stripnet_mona_last_n_stages: int = 2,
+        stripnet_freeze: bool = True,
     ) -> None:
         super().__init__()
         self.embed_dim = self.DINO_DIM  # native 1024 (StripNet projects 512 -> 1024)
@@ -320,12 +321,19 @@ class AsymmetricEncoder(nn.Module):
             # 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,
-            )
+            if stripnet_freeze:
+                self._freeze(self.image_encoder.backbone)
+                LOGGER.info("StripNet backbone: frozen (Conv-MONA + projection trainable)")
+            else:
+                LOGGER.info("StripNet backbone: UNFROZEN — full fine-tune (use lower lr factor)")
+            if stripnet_mona_last_n_stages > 0:
+                inject_conv_mona_into_stripnet(
+                    self.image_encoder.backbone,
+                    bottleneck=mona_bottleneck,
+                    last_n_stages=stripnet_mona_last_n_stages,
+                )
+            else:
+                LOGGER.info("Conv-MONA disabled (stripnet_mona_last_n_stages=0)")
             LOGGER.info("StripNet backbone: shared encoder, projection 512 -> %d", self.DINO_DIM)
         elif shared_encoder:
             self.image_encoder = DINOv3ViT.from_pretrained(dino_web_path)

src/training/train_gtauav.py

@@ -93,7 +93,9 @@ class TrainConfigGTAUAV:
     # 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
+    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)
 
     # Training.
     resume_from: str | None = None  # path to checkpoint for resuming
@@ -168,22 +170,37 @@ def _build_param_groups(
     model: AsymmetricEncoder,
     lr: float,
     text_lr_factor: float,
+    stripnet_backbone_lr_factor: float = 0.1,
 ) -> list[dict]:
-    """Build optimizer param groups with separate LR for text encoder."""
+    """Build optimizer param groups with separate LR for text encoder and unfrozen StripNet backbone.
+
+    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
+    """
     text_params = []
+    backbone_params = []
     other_params = []
 
+    is_stripnet = isinstance(getattr(model, "image_encoder", None), nn.Module) and \
+                  getattr(model, "backbone", "dinov3") == "stripnet"
+
     for name, param in model.named_parameters():
         if not param.requires_grad:
             continue
         if "text_encoder" in name:
             text_params.append(param)
+        elif is_stripnet and name.startswith("image_encoder.backbone.") and "mona_" not in name:
+            backbone_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})
+    if backbone_params:
+        groups.append({"params": backbone_params, "lr": lr * stripnet_backbone_lr_factor})
 
     return groups
 
@@ -586,6 +603,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
             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)
 
@@ -756,7 +774,10 @@ def train(cfg: TrainConfigGTAUAV) -> None:
     )
 
     # Optimizer — per-group LR (text encoder gets lower LR).
-    param_groups = _build_param_groups(model, cfg.learning_rate, cfg.text_lr_factor)
+    param_groups = _build_param_groups(
+        model, cfg.learning_rate, cfg.text_lr_factor,
+        stripnet_backbone_lr_factor=cfg.stripnet_backbone_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)