diff --git a/conf/gtauav_balanced_stripnet_unfrozen.gin b/conf/gtauav_balanced_stripnet_unfrozen.gin new file mode 100644 index 0000000..5e3a79e --- /dev/null +++ b/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" diff --git a/conf/gtauav_baseline_stripnet_unfrozen.gin b/conf/gtauav_baseline_stripnet_unfrozen.gin new file mode 100644 index 0000000..8d7df74 --- /dev/null +++ b/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 diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 872165e..123f2d9 100644 --- a/src/models/asymmetric_encoder.py +++ b/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) diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 7c51588..24ae3ea 100644 --- a/src/training/train_gtauav.py +++ b/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)