File size: 2,350 Bytes
c5a246d | 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 | from torch import nn
from transformers import PreTrainedModel
from transformers.models.dinov2.modeling_dinov2 import Dinov2Encoder
from .configuration_glint import AlignLayerConfig
def build_align_layer(config):
if config.model_type == "align_transformer":
model = AlignTransformer(config)
elif config.model_type == "identity":
model = Identity()
elif config.model_type == "mlp":
model = MLP(config)
elif config.model_type == "linear":
model = Linear(config)
else:
raise NotImplementedError()
return model
class AlignTransformer(PreTrainedModel):
def __init__(self, config: AlignLayerConfig):
super().__init__(config)
self.linear = nn.Linear(config.input_dim, config.hidden_size)
if config.num_hidden_layers:
self.transformer_layers = Dinov2Encoder(config)
else:
self.transformer_layers = None
if config.use_layer_norm:
self.layer_norm = nn.LayerNorm(config.hidden_size)
else:
self.layer_norm = None
def forward(self, vision_tokens):
if self.transformer_layers is not None:
vision_tokens = self.linear(vision_tokens)
vision_tokens = self.transformer_layers(vision_tokens)["last_hidden_state"]
if self.layer_norm is not None:
vision_tokens = self.layer_norm(vision_tokens)
return vision_tokens
class Identity(nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return x
class Linear(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.linear = nn.Linear(config.input_dim, config.hidden_size)
def forward(self, x):
return self.linear(x)
class MLP(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.mlp_layer = nn.Sequential(
nn.Linear(config.input_dim, config.input_dim),
nn.ReLU(inplace=True),
nn.Dropout(config.dropout),
nn.Linear(config.input_dim, config.input_dim),
nn.ReLU(inplace=True),
nn.Dropout(config.dropout),
nn.Linear(config.input_dim, config.hidden_size),
)
def forward(self, x):
return self.mlp_layer(x)
|