| |
| import torch |
| import torch.nn as nn |
| from transformers import PreTrainedModel |
| from transformers.activations import ACT2FN |
|
|
| from .configuration_projector import ProjectorConfig |
|
|
|
|
| class ProjectorModel(PreTrainedModel): |
| _auto_class = 'AutoModel' |
| config_class = ProjectorConfig |
| base_model_prefix = 'model' |
| supports_gradient_checkpointing = True |
|
|
| def __init__(self, config: ProjectorConfig) -> None: |
| super().__init__(config) |
| self.gradient_checkpointing = False |
|
|
| modules = [ |
| nn.Linear( |
| config.visual_hidden_size, |
| config.llm_hidden_size, |
| bias=config.bias) |
| ] |
| for _ in range(1, config.depth): |
| modules.append(ACT2FN[config.hidden_act]) |
| modules.append( |
| nn.Linear( |
| config.llm_hidden_size, |
| config.llm_hidden_size, |
| bias=config.bias)) |
| self.model = nn.Sequential(*modules) |
|
|
| def enable_input_require_grads(self): |
|
|
| def make_inputs_require_grad(module, input, output): |
| output.requires_grad_(True) |
|
|
| self.model.register_forward_hook(make_inputs_require_grad) |
|
|
| def _set_gradient_checkpointing(self, module, value=False): |
| if isinstance(module, ProjectorModel): |
| module.gradient_checkpointing = value |
|
|
| def forward(self, x): |
| if self.gradient_checkpointing and self.training: |
| layer_outputs = torch.utils.checkpoint.checkpoint(self.model, x) |
| else: |
| layer_outputs = self.model(x) |
| return layer_outputs |
|
|
