| 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) |
|
|
