File size: 5,317 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 | # 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 math
from dataclasses import dataclass
from typing import Optional, Tuple
import torch
import torch.nn.functional as F
from torch import nn
from dinov3.eval.text.text_tower import build_text_model
from dinov3.eval.text.vision_tower import build_vision_model
@dataclass
class DINOTxtConfig:
embed_dim: int
vision_backbone_config: str | None = None
text_backbone_config: str | None = None
vision_backbone_pretrained_weights: str | None = None
text_backbone_pretrained_weights: str | None = None
vision_model_freeze_backbone: bool = True
vision_model_train_img_size: int = 224
vision_model_use_class_token: bool = True
vision_model_use_patch_tokens: bool = False
vision_model_num_head_blocks: int = 0
vision_model_head_blocks_drop_path: float = 0.3
vision_model_use_linear_projection: bool = False
vision_model_patch_tokens_pooler_type: str = "mean"
vision_model_patch_token_layer: int = 1 # which layer to take patch tokens from
# 1 - last layer, 2 - second last layer, etc.
text_model_freeze_backbone: bool = False
text_model_num_head_blocks: int = 0
text_model_head_blocks_is_causal: bool = False
text_model_head_blocks_drop_prob: float = 0.0
text_model_tokens_pooler_type: str = "first"
text_model_use_linear_projection: bool = False
text_vocab_path_or_url: Optional[str] = None
init_logit_scale: float = math.log(1 / 0.07)
init_logit_bias: Optional[float] = None
freeze_logit_scale: bool = False
class DINOTxt(nn.Module):
def __init__(
self,
model_config: DINOTxtConfig,
vision_backbone: Optional[nn.Module] = None,
text_backbone: Optional[nn.Module] = None,
device=None,
):
super().__init__()
self.model_config = model_config
self.visual_model = build_vision_model(
model_config.embed_dim,
model_config.vision_backbone_config,
model_config.vision_model_freeze_backbone,
model_config.vision_model_num_head_blocks,
model_config.vision_model_head_blocks_drop_path,
model_config.vision_model_use_class_token,
model_config.vision_model_use_patch_tokens,
model_config.vision_model_patch_token_layer,
model_config.vision_model_patch_tokens_pooler_type,
model_config.vision_model_use_linear_projection,
backbone=vision_backbone,
)
self.text_model = build_text_model(
model_config.embed_dim,
model_config.text_backbone_config,
model_config.text_model_freeze_backbone,
model_config.text_model_num_head_blocks,
model_config.text_model_head_blocks_is_causal,
model_config.text_model_head_blocks_drop_prob,
model_config.text_model_tokens_pooler_type,
model_config.text_model_use_linear_projection,
backbone=text_backbone,
)
self.logit_scale = nn.Parameter(torch.empty(1, device=device))
if model_config.freeze_logit_scale:
self.logit_scale.requires_grad = False
def init_weights(self):
torch.nn.init.constant(self.logit_scale, self.model_config.init_logit_scale)
self.visual_model.init_weights()
self.text_model.init_weights()
def encode_image_with_patch_tokens(
self,
image: torch.Tensor,
normalize: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
features, patch_tokens, backbone_patch_tokens = self.visual_model(image)
return (
F.normalize(features, dim=-1) if normalize else features,
patch_tokens,
backbone_patch_tokens,
)
def encode_image(
self,
image: torch.Tensor,
normalize: bool = False,
) -> torch.Tensor:
features, _, _ = self.visual_model(image)
return F.normalize(features, dim=-1) if normalize else features
def encode_text(self, text: torch.Tensor, normalize: bool = False) -> torch.Tensor:
features = self.text_model(text)
return F.normalize(features, dim=-1) if normalize else features
def get_logits(
self, image: torch.Tensor, text: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
text_features = self.encode_text(text, normalize=True)
image_features = self.encode_image(image, normalize=True)
image_logits = self.logit_scale.exp() * image_features @ text_features.T
text_logits = image_logits.T
return image_logits, text_logits
def forward(
self,
image: torch.Tensor,
text: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
text_features = self.encode_text(text, normalize=True)
image_features, patch_tokens, backbone_patch_tokens = (
self.encode_image_with_patch_tokens(image, normalize=True)
)
return (
image_features,
text_features,
self.logit_scale.exp(),
patch_tokens,
backbone_patch_tokens,
)
|