initial commit

src/models/stripnet/conv_mona.py

@@ -0,0 +1,106 @@
+"""Conv-MONA: 2D adaptation of MONA (CVPR 2025) for hierarchical CNN backbones.
+
+MONA paper applies sequence-form adapters after MSA / MLP in ViT blocks. Here we
+mirror that idea in [B, C, H, W] form: BN → 1×1 Down(C→bn) → multi-scale DWConv
+{3,5,7} mean → +residual → GELU → 1×1 Up(bn→C). Layer scale (γ) channel-wise,
+init 1e-6 for near-identity start. Two adapters per StripNet Block: post-attn
+and post-mlp.
+"""
+from __future__ import annotations
+
+import logging
+import torch
+import torch.nn as nn
+
+from src.models.stripnet.model import StripNet, Block
+
+LOGGER = logging.getLogger("caption_test.stripnet.adapters")
+
+
+class ConvMona(nn.Module):
+    """Single Conv-MONA adapter.
+
+    Args:
+        dim: input channel dim.
+        bottleneck: bottleneck channel dim (e.g. 64).
+        gamma_init: layer-scale init value (1e-6 for near-identity at start).
+    """
+
+    def __init__(self, dim: int, bottleneck: int = 64, gamma_init: float = 1e-6) -> None:
+        super().__init__()
+        self.norm = nn.BatchNorm2d(dim)
+        self.down = nn.Conv2d(dim, bottleneck, kernel_size=1, bias=True)
+        self.dw3 = nn.Conv2d(bottleneck, bottleneck, kernel_size=3, padding=1, groups=bottleneck, bias=True)
+        self.dw5 = nn.Conv2d(bottleneck, bottleneck, kernel_size=5, padding=2, groups=bottleneck, bias=True)
+        self.dw7 = nn.Conv2d(bottleneck, bottleneck, kernel_size=7, padding=3, groups=bottleneck, bias=True)
+        self.act = nn.GELU()
+        self.up = nn.Conv2d(bottleneck, dim, kernel_size=1, bias=True)
+        # Channel-wise layer scale (γ), broadcast across H, W.
+        self.gamma = nn.Parameter(gamma_init * torch.ones(dim), requires_grad=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        h = self.down(h)
+        h = (self.dw3(h) + self.dw5(h) + self.dw7(h)) / 3.0 + h
+        h = self.act(h)
+        h = self.up(h)
+        return self.gamma.view(1, -1, 1, 1) * h
+
+
+def _patched_block_forward(block: Block, mona_attn: ConvMona, mona_mlp: ConvMona):
+    """Closure that wraps a Block.forward with two Conv-MONA residuals."""
+    orig_attn = block.attn
+    orig_mlp = block.mlp
+    orig_norm1 = block.norm1
+    orig_norm2 = block.norm2
+    orig_drop = block.drop_path
+    ls1 = block.layer_scale_1
+    ls2 = block.layer_scale_2
+
+    def forward(x: torch.Tensor) -> torch.Tensor:
+        x = x + orig_drop(ls1.unsqueeze(-1).unsqueeze(-1) * orig_attn(orig_norm1(x))) + mona_attn(x)
+        x = x + orig_drop(ls2.unsqueeze(-1).unsqueeze(-1) * orig_mlp(orig_norm2(x))) + mona_mlp(x)
+        return x
+
+    return forward
+
+
+def inject_conv_mona_into_stripnet(
+    model: StripNet,
+    bottleneck: int = 64,
+    last_n_stages: int = 2,
+    use_bf16: bool = False,
+) -> int:
+    """Inject Conv-MONA adapters into the deepest `last_n_stages` of StripNet.
+
+    Each Block in the targeted stages gets two adapters (post-attn, post-mlp).
+    Returns the number of adapters injected.
+
+    Stages are 1-indexed in StripNet (block1..block4). With `last_n_stages=2`
+    we adapt block3 and block4 — the deepest, semantically richest features.
+    """
+    n_stages = model.num_stages
+    target_stages = list(range(max(1, n_stages - last_n_stages + 1), n_stages + 1))
+    n_added = 0
+
+    for stage_idx in target_stages:
+        blocks: nn.ModuleList = getattr(model, f"block{stage_idx}")
+        dim = model.embed_dims[stage_idx - 1]
+        for blk_idx, block in enumerate(blocks):
+            mona_a = ConvMona(dim=dim, bottleneck=bottleneck)
+            mona_m = ConvMona(dim=dim, bottleneck=bottleneck)
+            if use_bf16:
+                mona_a.to(dtype=torch.bfloat16)
+                mona_m.to(dtype=torch.bfloat16)
+            # Register as submodules so they get moved by .to(device) / .train() etc.
+            block.add_module(f"mona_attn", mona_a)
+            block.add_module(f"mona_mlp", mona_m)
+            block.forward = _patched_block_forward(block, mona_a, mona_m)
+            n_added += 2
+
+    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    LOGGER.info(
+        "🔧 Conv-MONA injected: %d adapters in stages %s, %d trainable params (bottleneck=%d)",
+        n_added, target_stages, n_trainable, bottleneck,
+    )
+    return n_added