diff --git a/src/models/adapters.py b/src/models/adapters.py new file mode 100644 index 0000000..bd1cd1c --- /dev/null +++ b/src/models/adapters.py @@ -0,0 +1,286 @@ +from __future__ import annotations + +"""MONA and LoRA adapters for DINOv3 and DGTRS-CLIP. + +MONA (Multi-Cognitive One-Shot Nested Adaptation, CVPR 2025): + Vision-specific adapter with multi-scale depthwise convolutions. + Applied to DINOv3 ViT-L/16 (after MSA and MLP in each block). + Original: github.com/Leiyi-Hu/mona (MIT license). + +LoRA (Low-Rank Adaptation): + Low-rank matrices on Q/V attention projections. + Applied to DGTRS-CLIP text encoder (all 12 transformer blocks). +""" + +import logging +import math + +import coloredlogs +import torch +import torch.nn as nn +import torch.nn.functional as F + +LOGGER = logging.getLogger("caption_test.adapters") +coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s") + + +# --------------------------------------------------------------------------- +# MONA adapter (for DINOv3 ViT — 2D vision tokens) +# --------------------------------------------------------------------------- + +class MonaOp(nn.Module): + """Multi-cognitive visual filter: parallel depthwise convs (3×3, 5×5, 7×7).""" + + def __init__(self, channels: int) -> None: + super().__init__() + self.conv3 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, groups=channels) + self.conv5 = nn.Conv2d(channels, channels, kernel_size=5, padding=2, groups=channels) + self.conv7 = nn.Conv2d(channels, channels, kernel_size=7, padding=3, groups=channels) + self.projector = nn.Conv2d(channels, channels, kernel_size=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + identity = x + x = (self.conv3(x) + self.conv5(x) + self.conv7(x)) / 3.0 + identity + return x + self.projector(x) + + +class MonaAdapter(nn.Module): + """MONA adapter module for ViT blocks. + + ScaledLayerNorm → Down(dim→bottleneck) → reshape 2D → MonaOp → reshape 1D + → GELU → Dropout → Up(bottleneck→dim) + residual + + Args: + dim: Input feature dimension (e.g. 1024 for ViT-L). + bottleneck: Bottleneck dimension for down/up projections. + dropout: Dropout rate. + """ + + def __init__(self, dim: int = 1024, bottleneck: int = 64, dropout: float = 0.1) -> None: + super().__init__() + self.down = nn.Linear(dim, bottleneck) + self.up = nn.Linear(bottleneck, dim) + self.mona_op = MonaOp(bottleneck) + self.norm = nn.LayerNorm(dim) + self.dropout = nn.Dropout(p=dropout) + # Scaled init: gamma ≈ 0 at start (near-identity), gammax ≈ 1. + self.gamma = nn.Parameter(torch.ones(dim) * 1e-6) + self.gammax = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor, hw: tuple[int, int]) -> torch.Tensor: + """Apply MONA adapter (runs in fp32 to avoid AMP overflow). + + Args: + x: Token features [B, N, D] where N includes CLS + register + patch tokens. + hw: (H, W) spatial shape of patch tokens. + + Returns: + Adapted features [B, N, D] (same shape, residual connection). + """ + with torch.amp.autocast("cuda", enabled=False): + return self._forward_fp32(x.float(), hw) + + def _forward_fp32(self, x: torch.Tensor, hw: tuple[int, int]) -> torch.Tensor: + identity = x + # Scaled LayerNorm. + x = self.norm(x) * self.gamma + x * self.gammax + + x = self.down(x) # [B, N, bottleneck] + + B, N, C = x.shape + H, W = hw + n_special = N - H * W # CLS + register tokens + + # Separate special tokens (CLS, registers) from patch tokens. + special = x[:, :n_special] # [B, n_special, C] + patches = x[:, n_special:] # [B, H*W, C] + + # Reshape patches to 2D for convolutions. + patches = patches.reshape(B, H, W, C).permute(0, 3, 1, 2) # [B, C, H, W] + patches = self.mona_op(patches) + patches = patches.permute(0, 2, 3, 1).reshape(B, H * W, C) # [B, H*W, C] + + # Recombine. + x = torch.cat([special, patches], dim=1) # [B, N, C] + + x = F.gelu(x) + x = self.dropout(x) + x = self.up(x) # [B, N, D] + + return identity + x + + +# --------------------------------------------------------------------------- +# LoRA adapter (for DGTRS-CLIP text encoder — 1D text tokens) +# --------------------------------------------------------------------------- + +class LoRALinear(nn.Module): + """LoRA adapter for a linear layer: output = W·x + (B·A)·x. + + Args: + in_features: Input dimension. + out_features: Output dimension. + rank: LoRA rank (low-rank decomposition dimension). + alpha: LoRA scaling factor (effective scale = alpha / rank). + """ + + def __init__( + self, + in_features: int, + out_features: int, + rank: int = 4, + alpha: float = 1.0, + ) -> None: + super().__init__() + self.rank = rank + self.scale = alpha / rank + # A: down-projection, B: up-projection. Init: A ~ N(0, 1/rank), B = 0. + self.A = nn.Parameter(torch.randn(rank, in_features) / math.sqrt(rank)) + self.B = nn.Parameter(torch.zeros(out_features, rank)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Compute LoRA delta: (B @ A) @ x * scale. Always fp32.""" + with torch.amp.autocast("cuda", enabled=False): + x = x.float() + return (x @ self.A.t() @ self.B.t()) * self.scale + + +# --------------------------------------------------------------------------- +# Injection functions — apply adapters to existing frozen models +# --------------------------------------------------------------------------- + +def inject_mona_into_dinov3( + model: nn.Module, + bottleneck: int = 64, + dropout: float = 0.1, +) -> int: + """Inject MONA adapters into a frozen DINOv3ViT model. + + Adds two MonaAdapter modules per block (after attention, after MLP). + Only adapter parameters are trainable. + + Args: + model: DINOv3ViT model (already frozen). + bottleneck: MONA bottleneck dimension. + dropout: MONA dropout rate. + + Returns: + Number of trainable adapter parameters added. + """ + dim = model.embed_dim + n_adapters = 0 + + for block in model.layer: + # Create adapters. + block.mona_attn = MonaAdapter(dim, bottleneck, dropout) + block.mona_mlp = MonaAdapter(dim, bottleneck, dropout) + + # Wrap the original forward to inject adapters. + original_forward = block.forward + + def make_mona_forward(blk, orig_fwd): + def mona_forward(x: torch.Tensor) -> torch.Tensor: + B, N, C = x.shape + # DINOv3: [CLS, 4 registers, H*W patches] + n_special = 1 + 4 # CLS + registers + n_patches = N - n_special + H = W = int(math.sqrt(n_patches)) + + # Original block forward (attention + MLP with layer scale). + x = x + blk.layer_scale1(blk.attention(blk.norm1(x))) + x = blk.mona_attn(x, (H, W)) + x = x + blk.layer_scale2(blk.mlp(blk.norm2(x))) + x = blk.mona_mlp(x, (H, W)) + return x + return mona_forward + + block.forward = make_mona_forward(block, original_forward) + n_adapters += 2 + + n_params = sum(p.numel() for n, p in model.named_parameters() if "mona" in n) + LOGGER.info( + "🔧 MONA injected: %d adapters (%d per block), %s trainable params (bottleneck=%d)", + n_adapters, 2, f"{n_params:,}", bottleneck, + ) + return n_params + + +def inject_lora_into_dgtrs( + model: nn.Module, + rank: int = 4, + alpha: float = 1.0, +) -> int: + """Inject LoRA adapters into DGTRS-CLIP text encoder attention layers. + + Adds LoRA to Q and V projections in each ResidualAttentionBlock. + Original weights stay frozen, LoRA params are trainable. + + Args: + model: DGTRSTextEncoder model. + rank: LoRA rank. + alpha: LoRA scaling factor. + + Returns: + Number of trainable LoRA parameters added. + """ + n_adapted = 0 + + for i, block in enumerate(model.transformer.resblocks): + attn = block.attn + # nn.MultiheadAttention uses in_proj_weight [3*dim, dim] for Q, K, V. + embed_dim = attn.embed_dim + + # LoRA on Q and V (not K — standard practice). + block.lora_q = LoRALinear(embed_dim, embed_dim, rank, alpha) + block.lora_v = LoRALinear(embed_dim, embed_dim, rank, alpha) + + # Wrap attention forward to add LoRA deltas. + original_attention = block.attention + + def make_lora_attention(blk): + def lora_attention(x: torch.Tensor) -> torch.Tensor: + # x shape: [T, B, D] (sequence-first for nn.MultiheadAttention) + attn_mod = blk.attn + attn_mask = blk.attn_mask + if attn_mask is not None: + attn_mask = attn_mask.to(dtype=x.dtype, device=x.device) + + # Compute QKV using original in_proj_weight. + T, B, D = x.shape + + # Original QKV. + qkv = F.linear(x, attn_mod.in_proj_weight, attn_mod.in_proj_bias) + q, k, v = qkv.chunk(3, dim=-1) + + # Add LoRA deltas to Q and V. + # LoRA expects [B, T, D], but x is [T, B, D]. + x_btd = x.permute(1, 0, 2) # [B, T, D] + q = q + blk.lora_q(x_btd).permute(1, 0, 2) + v = v + blk.lora_v(x_btd).permute(1, 0, 2) + + # Multi-head attention. + num_heads = attn_mod.num_heads + head_dim = D // num_heads + q = q.reshape(T, B * num_heads, head_dim).transpose(0, 1) + k = k.reshape(T, B * num_heads, head_dim).transpose(0, 1) + v = v.reshape(T, B * num_heads, head_dim).transpose(0, 1) + + attn_output = F.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, + ) + attn_output = attn_output.transpose(0, 1).reshape(T, B, D) + attn_output = attn_mod.out_proj(attn_output) + return attn_output + return lora_attention + + block.attention = make_lora_attention(block) + n_adapted += 1 + + n_params = sum( + p.numel() for n, p in model.named_parameters() if "lora" in n + ) + LOGGER.info( + "🔧 LoRA injected: %d blocks, rank=%d, %s trainable params", + n_adapted, rank, f"{n_params:,}", + ) + return n_params diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 9ca6dd1..dac6c0c 100644 --- a/src/models/asymmetric_encoder.py +++ b/src/models/asymmetric_encoder.py @@ -25,6 +25,7 @@ LOGGER = logging.getLogger("caption_test.model") coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s") 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 @@ -279,6 +280,8 @@ class AsymmetricEncoder(nn.Module): lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt", init_gate: float = 0.7, baseline_mode: bool = False, + mona_bottleneck: int = 64, + lora_rank: int = 4, device: str = "cuda", ) -> None: super().__init__() @@ -286,17 +289,19 @@ class AsymmetricEncoder(nn.Module): self.baseline_mode = baseline_mode self.device = device - # Image encoders (frozen). + # Image encoders (frozen + MONA adapters). self.drone_encoder = DINOv3ViT.from_pretrained(dino_web_path) self.sat_encoder = DINOv3ViT.from_pretrained(dino_sat_path) self._freeze(self.drone_encoder) self._freeze(self.sat_encoder) + inject_mona_into_dinov3(self.drone_encoder, bottleneck=mona_bottleneck) + inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck) - # Text encoder — official DGTRS architecture (partial unfreeze). + # Text encoder — official DGTRS architecture (frozen + LoRA). if not baseline_mode: self.text_encoder = load_dgtrs_text_encoder(lrsclip_path) self._freeze(self.text_encoder) - self._unfreeze_text_last_block() + inject_lora_into_dgtrs(self.text_encoder, rank=lora_rank) else: self.text_encoder = None @@ -317,33 +322,13 @@ class AsymmetricEncoder(nn.Module): p.requires_grad = False module.eval() - def _unfreeze_text_last_block(self) -> None: - """Unfreeze last transformer block + text_projection + ln_final.""" - if self.text_encoder is None: - return - # Last resblock. - for p in self.text_encoder.transformer.resblocks[-1].parameters(): - p.requires_grad = True - # ln_final. - for p in self.text_encoder.ln_final.parameters(): - p.requires_grad = True - # text_projection. - tp = self.text_encoder.text_projection - if isinstance(tp, nn.Parameter): - tp.requires_grad = True - elif isinstance(tp, nn.Module): - for p in tp.parameters(): - p.requires_grad = True - def encode_drone(self, images: torch.Tensor) -> torch.Tensor: - """Encode drone images. Returns [B, DINO_DIM].""" - with torch.no_grad(): - return self.drone_encoder(images) + """Encode drone images with MONA adapters. Returns [B, DINO_DIM].""" + return self.drone_encoder(images) def encode_satellite(self, images: torch.Tensor) -> torch.Tensor: - """Encode satellite images. Returns [B, DINO_DIM].""" - with torch.no_grad(): - return self.sat_encoder(images) + """Encode satellite images with MONA adapters. Returns [B, DINO_DIM].""" + return self.sat_encoder(images) def encode_text_levels( self,