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>
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user