Add MONA adapters for DINOv3 + LoRA for DGTRS-CLIP text encoder
MONA (Multi-Cognitive One-Shot Nested Adaptation, CVPR 2025): - 2 adapters per DINOv3 block (after MSA, after MLP) × 24 layers × 2 encoders - MonaOp: parallel DWConv 3×3/5×5/7×7 + 1×1 projector - ScaledLayerNorm + down(1024→64) + MonaOp + GELU + up(64→1024) - 7M params per encoder, 14M total for drone+sat LoRA (Low-Rank Adaptation): - Q and V projections in all 12 DGTRS-CLIP transformer blocks - rank=4, ~147K params total - Replaces partial unfreeze (was ~7.6M for last block) Both adapters run in fp32 (torch.amp.autocast disabled) to avoid AMP gradient overflow. Frozen backbone layers still run in fp16. Total trainable: 17.6M / 748M (2.35%) MONA (2×DINOv3): 14.0M LoRA (DGTRS): 147K TextFusionMLP: 3.4M Gates + logit: 3 Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
286
src/models/adapters.py
Normal file
286
src/models/adapters.py
Normal file
@@ -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
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user