| |
| import torch |
| import torch.nn as nn |
| from transformers import PreTrainedModel |
| from transformers.activations import ACT2FN |
|
|
| from .configuration_downsampler import DownsamplerConfig |
|
|
|
|
| class DownsamplerModel(PreTrainedModel): |
| _auto_class = 'AutoModel' |
| config_class = DownsamplerConfig |
| base_model_prefix = 'model' |
| supports_gradient_checkpointing = True |
|
|
| def __init__(self, config: DownsamplerConfig) -> None: |
| super().__init__(config) |
| self.gradient_checkpointing = False |
|
|
| self.group_op = nn.Conv2d( |
| in_channels=config.visual_hidden_size, |
| out_channels=config.llm_hidden_size, |
| bias=config.bias, |
| kernel_size=config.kernel_size, stride=config.stride) |
| modules = list() |
| 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.linear_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, DownsamplerModel): |
| module.gradient_checkpointing = value |
|
|
| def _forward(self, x): |
|
|
| |
| x = x.permute(0, 3, 1, 2) |
| x = self.group_op(x) |
| |
| x = x.permute(0, 2, 3, 1) |
| x = self.linear_model(x) |
|
|
| return x |
|
|
| def forward(self, x): |
| if self.gradient_checkpointing and self.training: |
| layer_outputs = torch.utils.checkpoint.checkpoint(self._forward, x) |
| else: |
| layer_outputs = self._forward(x) |
| return layer_outputs |
|
|
