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:
pikaliov
2026-04-25 16:04:24 +03:00
parent d4cb2dd300
commit c6fcd2222c
4 changed files with 63 additions and 9 deletions

View File

@@ -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)