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