| import torch |
| torch.manual_seed(1024) |
|
|
| import torch.nn as nn |
| from transformers import PreTrainedModel |
|
|
| from .configuration_hformer import HformerConfig |
| from .qformer_src import BertConfig, BertLMHeadModel |
|
|
| from transformers import BertTokenizerFast as BertTokenizer |
|
|
| from .configuration_projector import ProjectorConfig |
| from .modeling_projector import ProjectorModel |
| import torch.nn.functional as F |
| from transformers.activations import ACT2FN |
|
|
|
|
| class LayerNorm(nn.LayerNorm): |
| def forward(self, x: torch.Tensor): |
| ret = super().forward(x) |
| return ret |
|
|
| class HformerModel(PreTrainedModel): |
| _auto_class = 'AutoModel' |
| config_class = HformerConfig |
| base_model_prefix = 'model' |
| supports_gradient_checkpointing = False |
|
|
| def __init__(self, config) -> None: |
| super().__init__(config) |
| self.gradient_checkpointing = False |
| vision_width = config.visual_hidden_size |
| num_query_token = config.num_query_token |
| bert = config.bert |
| llm_hidden_size = config.llm_hidden_size |
| cross_attention_freq = config.cross_attention_freq |
| qformer_pth = config.qformer_pth |
|
|
| encoder_config = BertConfig.from_pretrained(bert) |
| encoder_config.encoder_width = vision_width |
| encoder_config.add_cross_attention = True |
| encoder_config.cross_attention_freq = cross_attention_freq |
| encoder_config.query_length = num_query_token |
| encoder_config.num_hidden_layers = 12 |
| Qformer = BertLMHeadModel.from_pretrained( |
| bert, config=encoder_config |
| ) |
| remove_text = False |
| if remove_text: |
| Qformer.cls = None |
| Qformer.bert.embeddings.word_embeddings = None |
| Qformer.bert.embeddings.position_embeddings = None |
| for layer in Qformer.bert.encoder.layer: |
| layer.output = None |
| layer.intermediate = None |
|
|
| query_tokens = nn.Parameter( |
| torch.zeros(1, num_query_token, encoder_config.hidden_size) |
| ) |
| query_tokens.data.normal_(mean=0.0, std=encoder_config.initializer_range) |
| |
| self.Qformer = Qformer |
| self.query_tokens = query_tokens |
| self.llm_proj = nn.Linear(encoder_config.hidden_size, llm_hidden_size, bias=config.bias) |
| self.ln_vision = LayerNorm(encoder_config.encoder_width) |
| self.ln_llava = LayerNorm(encoder_config.encoder_width) |
| |
| tokenizer = BertTokenizer.from_pretrained(bert, truncation_side='right') |
| tokenizer.add_special_tokens({"bos_token": "[DEC]"}) |
| self.Qformer.resize_token_embeddings(len(tokenizer)) |
|
|
| if qformer_pth is not None: |
| pretrained_state_dict = torch.load(qformer_pth, map_location='cpu')['model'] |
| print(f'Load Qformer from {qformer_pth}') |
| self.load_state_dict(pretrained_state_dict, strict=False) |
| print('Done.') |
|
|
| projector_config = ProjectorConfig( |
| visual_hidden_size = config.visual_hidden_size, |
| llm_hidden_size = config.llm_hidden_size, |
| projector_depth = 2) |
| self.connector = ProjectorModel(projector_config) |
|
|
| modules = [ |
| nn.Linear(config.llm_hidden_size, config.llm_hidden_size//4, bias=False), |
| ACT2FN['gelu'], |
| nn.Linear(config.llm_hidden_size//4, config.llm_hidden_size, bias=False) |
| ] |
| self.ffn = nn.Sequential(*modules) |
|
|
| def enable_input_require_grads(self): |
| def make_inputs_require_grad(module, input, output): |
| if isinstance(output, tuple): |
| output[0].requires_grad_(True) |
| output[1].requires_grad_(True) |
| else: |
| output.requires_grad_(True) |
|
|
| self.Qformer.register_forward_hook(make_inputs_require_grad) |
| self.llm_proj.register_forward_hook(make_inputs_require_grad) |
| self.ln_vision.register_forward_hook(make_inputs_require_grad) |
| self.connector.register_forward_hook(make_inputs_require_grad) |
| self.ffn.register_forward_hook(make_inputs_require_grad) |
|
|
| def _set_gradient_checkpointing(self, module, value=False): |
| pass |
|
|
| def forward(self, x_): |
| if self.gradient_checkpointing and self.training: |
| print('Not support gradient checkpointing') |
| x = self.ln_vision(x_) |
| query_tokens = self.query_tokens.expand(x.shape[0], -1, -1) |
| query_output = self.Qformer.bert( |
| query_embeds=query_tokens, |
| encoder_hidden_states=x, |
| return_dict=True, |
| ) |
| |
| q_feat = self.llm_proj(query_output.last_hidden_state) |
| |
| mlp_outputs = self.connector(x_) |
| mlp_feat = mlp_outputs |
|
|
| int_feat = mlp_feat + q_feat.mean(dim=1)[:,None] |
| out = int_feat + self.ffn(int_feat) |
|
|
| return out |
|
|
|
|
