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