107 lines
4.2 KiB
Python
107 lines
4.2 KiB
Python
"""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
|