Add StripNet backbone option with Conv-MONA adaptation
StripNet-small (Strip-R-CNN, HVision-NKU) as alternative image encoder to
DINOv3 ViT-L/16. ~28M params (10x smaller). Output 512-dim from stage 4
projected to 1024 to keep retrieval space compatible with DINOv3 configs.
- src/models/stripnet/: self-contained backbone (model.py, conv_mona.py).
State-dict naming follows upstream Strip-R-CNN repo (conv_spatial1/2);
ImageNet-1K pretrained head dropped on load.
- Conv-MONA: 2D adaptation of MONA (CVPR 2025) for CNN blocks. BN → 1x1
Down(C->bn) → multi-scale DWConv {3,5,7} mean → +residual → GELU →
1x1 Up(bn->C) with channel-wise layer scale γ init 1e-6. Two adapters
per StripNet Block (post-attn, post-mlp); injected into deepest N stages.
- StripNetEncoder: GAP + Linear(512->1024). Overrides train() to keep
frozen BatchNorm stats stable across mode flips.
- AsymmetricEncoder: new `backbone="stripnet"` option (always shared).
- TrainConfigGTAUAV: backbone, stripnet_path, stripnet_mona_last_n_stages.
- conf/gtauav_balanced_stripnet.gin + gtauav_baseline_stripnet.gin.
Smoke test: forward [2,3,256,256] -> [2,1024]. Trainable: 1.2M baseline
(8.27%), 4.76M with text (3.35% of 142M).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user