File size: 5,470 Bytes
d9bb75c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | # Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This software may be used and distributed in accordance with
# the terms of the DINOv3 License Agreement.
import logging
from typing import Optional
import torch
from dinov3.eval.text.text_transformer import TextTransformer
from dinov3.layers import CausalSelfAttentionBlock
from torch import nn
logger = logging.getLogger("dinov3")
class TextHead(nn.Module):
def __init__(
self,
input_dim: int,
embed_dim: int,
num_heads: int,
num_blocks: int,
block_drop_prob: float,
is_causal: bool,
use_linear_projection: bool,
):
super().__init__()
block_list = [nn.Identity()]
self.ln_final = nn.Identity()
if num_blocks > 0:
logger.info(f"Adding {num_blocks} text tower transformer head blocks")
block_list = [
CausalSelfAttentionBlock(
dim=input_dim,
num_heads=num_heads,
is_causal=is_causal,
dropout_prob=block_drop_prob,
)
for _ in range(num_blocks)
]
self.ln_final = nn.LayerNorm(input_dim)
self.blocks = nn.ModuleList(block_list)
self.num_blocks = num_blocks
self.linear_projection = nn.Identity()
if input_dim != embed_dim or use_linear_projection:
logger.info(
f"Text tower : Using a linear projection from {input_dim} to {embed_dim}"
)
self.linear_projection = nn.Linear(input_dim, embed_dim, bias=False)
def init_weights(self):
if self.num_blocks > 0:
for i in range(self.num_blocks):
self.blocks[i].init_weights()
self.ln_final.reset_parameters()
if isinstance(self.linear_projection, nn.Linear):
nn.init.normal_(
self.linear_projection.weight,
std=self.linear_projection.in_features**-0.5,
)
def forward(self, text_tokens: torch.Tensor) -> torch.Tensor:
for block in self.blocks:
text_tokens = block(text_tokens)
text_tokens = self.ln_final(text_tokens)
return self.linear_projection(text_tokens)
class TextTower(nn.Module):
def __init__(
self,
backbone: nn.Module,
freeze_backbone: bool,
embed_dim: int,
num_head_blocks: int,
head_blocks_is_causal: bool,
head_blocks_block_drop_prob: float,
tokens_pooler_type: str,
use_linear_projection: bool,
):
super().__init__()
self.backbone = backbone
self.freeze_backbone = freeze_backbone
backbone_out_dim = backbone.embed_dim
logger.info(f"Text backbone embedding dimension: {backbone_out_dim}")
self.backbone = backbone
self.head = TextHead(
backbone_out_dim,
embed_dim,
self.backbone.num_heads,
num_head_blocks,
head_blocks_block_drop_prob,
head_blocks_is_causal,
use_linear_projection,
)
self.tokens_pooler_type = tokens_pooler_type
def init_weights(self):
self.backbone.init_weights()
self.head.init_weights()
def forward(self, token_indices: torch.Tensor) -> torch.Tensor:
text_tokens = self.backbone(token_indices)
text_tokens = self.head(text_tokens)
if self.tokens_pooler_type == "first":
features = text_tokens[:, 0]
elif self.tokens_pooler_type == "last":
features = text_tokens[:, -1]
elif self.tokens_pooler_type == "argmax":
assert token_indices is not None
features = text_tokens[
torch.arange(text_tokens.shape[0]), token_indices.argmax(dim=-1)
]
else:
raise ValueError(f"Unknown text tokens pooler type: {self.pooler_type}")
return features
def build_text_backbone(
cfg,
) -> torch.nn.Module:
logger.info("Setting up a text transformer")
model = TextTransformer(
context_length=cfg.context_length,
vocab_size=cfg.vocab_size,
dim=cfg.dim,
num_heads=cfg.num_heads,
num_layers=cfg.num_layers,
ffn_ratio=cfg.ffn_ratio,
is_causal=cfg.is_causal,
ls_init_value=cfg.ls_init_value,
dropout_prob=cfg.dropout_prob,
)
logger.info(f"Setting upa custom text transformer {cfg.model_name}")
return model
def build_text_model(
embed_dim: int,
backbone_model_config: str,
freeze_backbone: bool,
num_head_blocks: int,
head_blocks_is_causal: bool,
head_blocks_drop_prob: float,
tokens_pooler_type: str,
use_linear_projection: bool,
backbone: Optional[nn.Module] = None,
):
if backbone is None:
if backbone_model_config is not None:
from omegaconf import OmegaConf
cfg = OmegaConf.load(backbone_model_config)
backbone = build_text_backbone(cfg)
else:
raise RuntimeError(
"Failed to create, text backbone, either backbone or backbone_model_config should be not None"
)
return TextTower(
backbone,
freeze_backbone,
embed_dim,
num_head_blocks,
head_blocks_is_causal,
head_blocks_drop_prob,
tokens_pooler_type,
use_linear_projection,
)
|