Replace broken LRSCLIPTextEncoder with official DGTRS architecture
Root cause of NaN: our open_clip wrapper had 3 bugs: 1. Positional embeddings summed for all positions instead of masked (official: mask1 for pos 0-19, mask2 for pos 20-247) 2. open_clip uses batch-first transformer, DGTRS uses sequence-first (LND format with nn.MultiheadAttention) 3. open_clip tokenizer truncates to 77 tokens, DGTRS needs 248 Fix: copied official DGTRS text encoder architecture from github.com/MitsuiChen14/DGTRS (Apache-2.0): - src/models/dgtrs/model.py: DGTRSTextEncoder, build_model, load_dgtrs_text_encoder, tokenize_dgtrs - src/models/dgtrs/simple_tokenizer.py: BPE tokenizer (248 tokens) - src/models/dgtrs/bpe_simple_vocab_16e6.txt.gz: vocabulary Removed: LRSCLIPTextEncoder class, open_clip dependency for text encoding Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
@@ -17,7 +17,6 @@ import warnings
|
||||
from pathlib import Path
|
||||
|
||||
import coloredlogs
|
||||
import open_clip
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
@@ -26,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.dgtrs.model import DGTRSTextEncoder, load_dgtrs_text_encoder, tokenize_dgtrs
|
||||
from src.models.dual_encoder import GatedFusion, ProjectionHead
|
||||
|
||||
|
||||
@@ -197,116 +197,8 @@ class DINOv3ViT(nn.Module):
|
||||
return model
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# LRSCLIP (DGTRS-CLIP) text encoder — open_clip with KPS positional embedding
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class LRSCLIPTextEncoder(nn.Module):
|
||||
"""DGTRS-CLIP text encoder with KPS positional embedding (248 tokens).
|
||||
|
||||
Wraps open_clip ViT-L-14 text tower. Handles the extra
|
||||
positional_embedding_res (KPS residual) not present in vanilla open_clip.
|
||||
|
||||
Checkpoint keys (text-only):
|
||||
token_embedding.weight: [49408, 768]
|
||||
positional_embedding: [248, 768]
|
||||
positional_embedding_res: [248, 768]
|
||||
transformer.resblocks.{0-11}.{attn,mlp,ln}.*
|
||||
ln_final.{weight,bias}
|
||||
text_projection: [768, 768]
|
||||
"""
|
||||
|
||||
def __init__(self, context_length: int = 248, embed_dim: int = 768) -> None:
|
||||
super().__init__()
|
||||
self.context_length = context_length
|
||||
self.embed_dim = embed_dim
|
||||
|
||||
# Build the open_clip model to get correct architecture, then
|
||||
# replace positional embedding and add KPS residual.
|
||||
# Suppress "no pretrained weights" warning — we load DGTRS-CLIP weights after.
|
||||
_root = logging.getLogger()
|
||||
_prev_level = _root.level
|
||||
_root.setLevel(logging.ERROR)
|
||||
clip_model = open_clip.create_model("ViT-L-14", pretrained=None)
|
||||
_root.setLevel(_prev_level)
|
||||
|
||||
# Extract text components.
|
||||
self.token_embedding = clip_model.token_embedding
|
||||
self.transformer = clip_model.transformer
|
||||
self.ln_final = clip_model.ln_final
|
||||
self.text_projection = clip_model.text_projection
|
||||
self.attn_mask = None # rebuilt in forward
|
||||
|
||||
# Replace positional embedding with 248-length version.
|
||||
self.positional_embedding = nn.Parameter(
|
||||
torch.zeros(context_length, embed_dim),
|
||||
)
|
||||
self.positional_embedding_res = nn.Parameter(
|
||||
torch.zeros(context_length, embed_dim),
|
||||
)
|
||||
|
||||
def _build_attn_mask(self, context_length: int, device: torch.device) -> torch.Tensor:
|
||||
mask = torch.empty(context_length, context_length, device=device)
|
||||
mask.fill_(float("-inf"))
|
||||
mask.triu_(1)
|
||||
return mask
|
||||
|
||||
def forward(self, text: torch.Tensor) -> torch.Tensor:
|
||||
"""Encode tokenized text.
|
||||
|
||||
Args:
|
||||
text: Token IDs [B, T] (T <= 248).
|
||||
|
||||
Returns:
|
||||
Text embeddings [B, 768], L2-normalized.
|
||||
"""
|
||||
T = text.shape[1]
|
||||
x = self.token_embedding(text) # [B, T, 768]
|
||||
pos = self.positional_embedding[:T] + self.positional_embedding_res[:T]
|
||||
x = x + pos
|
||||
|
||||
if self.attn_mask is None or self.attn_mask.shape[0] != T:
|
||||
self.attn_mask = self._build_attn_mask(T, x.device)
|
||||
attn_mask = self.attn_mask[:T, :T]
|
||||
|
||||
# open_clip transformer expects batch-first [B, T, D].
|
||||
x = self.transformer(x, attn_mask=attn_mask)
|
||||
x = self.ln_final(x) # [B, T, D]
|
||||
|
||||
# Take features at EOS token position.
|
||||
eos_idx = text.argmax(dim=-1)
|
||||
x = x[torch.arange(x.shape[0], device=x.device), eos_idx]
|
||||
|
||||
# Project.
|
||||
if isinstance(self.text_projection, nn.Parameter):
|
||||
x = x @ self.text_projection
|
||||
else:
|
||||
x = self.text_projection(x)
|
||||
|
||||
return F.normalize(x, dim=-1)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls, path: str | Path) -> LRSCLIPTextEncoder:
|
||||
"""Load DGTRS-CLIP checkpoint, extracting only text encoder weights."""
|
||||
model = cls()
|
||||
path = Path(path)
|
||||
LOGGER.info("📝 Loading LRSCLIP text encoder from %s", path.name)
|
||||
full_state = torch.load(str(path), map_location="cpu", weights_only=False)
|
||||
if "state_dict" in full_state:
|
||||
full_state = full_state["state_dict"]
|
||||
|
||||
# Filter out visual.* keys — keep only text encoder.
|
||||
text_state = {
|
||||
k: v for k, v in full_state.items()
|
||||
if not k.startswith("visual.")
|
||||
}
|
||||
|
||||
# Handle text_projection shape: checkpoint may be [768, 768] Parameter
|
||||
# while open_clip stores it as nn.Parameter directly.
|
||||
model.load_state_dict(text_state, strict=False)
|
||||
n_params = sum(p.numel() for p in model.parameters())
|
||||
LOGGER.info("📝 LRSCLIP loaded: %s params, context=%d tokens", f"{n_params:,}", model.context_length)
|
||||
return model
|
||||
# LRSCLIPTextEncoder removed — replaced by official DGTRS architecture
|
||||
# in src/models/dgtrs/model.py (DGTRSTextEncoder)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -396,15 +288,13 @@ class AsymmetricEncoder(nn.Module):
|
||||
self._freeze(self.drone_encoder)
|
||||
self._freeze(self.sat_encoder)
|
||||
|
||||
# Text encoder (partial unfreeze).
|
||||
# Text encoder — official DGTRS architecture (partial unfreeze).
|
||||
if not baseline_mode:
|
||||
self.text_encoder = LRSCLIPTextEncoder.from_pretrained(lrsclip_path)
|
||||
self.text_encoder = load_dgtrs_text_encoder(lrsclip_path)
|
||||
self._freeze(self.text_encoder)
|
||||
self._unfreeze_text_last_block()
|
||||
self.tokenizer = open_clip.get_tokenizer("ViT-L-14")
|
||||
else:
|
||||
self.text_encoder = None
|
||||
self.tokenizer = None
|
||||
|
||||
# Projection heads.
|
||||
self.proj_drone = ProjectionHead(
|
||||
@@ -473,12 +363,8 @@ class AsymmetricEncoder(nn.Module):
|
||||
return F.normalize(fused, dim=-1)
|
||||
|
||||
def _encode_single_text(self, texts: list[str]) -> torch.Tensor:
|
||||
"""Tokenize and encode a list of strings."""
|
||||
tokens = self.tokenizer(list(texts)).to(self.device)
|
||||
# Pad/truncate to context_length.
|
||||
T = tokens.shape[1]
|
||||
if T > self.text_encoder.context_length:
|
||||
tokens = tokens[:, :self.text_encoder.context_length]
|
||||
"""Tokenize and encode a list of strings using DGTRS tokenizer."""
|
||||
tokens = tokenize_dgtrs(list(texts)).to(self.device)
|
||||
return self.text_encoder(tokens)
|
||||
|
||||
def forward(
|
||||
|
||||
Reference in New Issue
Block a user