| 355 | |
| 356 | |
| 357 | class ViTEncoder(nn.Module): |
| 358 | def __init__(self, config: ViTConfig) -> None: |
| 359 | super().__init__() |
| 360 | self.config = config |
| 361 | self.layer = nn.ModuleList([ViTLayer(config) for _ in range(config.num_hidden_layers)]) |
| 362 | self.gradient_checkpointing = False |
| 363 | |
| 364 | def forward( |
| 365 | self, |
| 366 | hidden_states: torch.Tensor, |
| 367 | adaln_input: torch.Tensor = None, |
| 368 | head_mask: Optional[torch.Tensor] = None, |
| 369 | output_attentions: bool = False, |
| 370 | output_hidden_states: bool = False, |
| 371 | return_dict: bool = True, |
| 372 | ) -> Union[tuple, BaseModelOutput]: |
| 373 | all_hidden_states = () if output_hidden_states else None |
| 374 | all_self_attentions = () if output_attentions else None |
| 375 | |
| 376 | for i, layer_module in enumerate(self.layer): |
| 377 | if output_hidden_states: |
| 378 | all_hidden_states = all_hidden_states + (hidden_states,) |
| 379 | |
| 380 | layer_head_mask = head_mask[i] if head_mask is not None else None |
| 381 | |
| 382 | if self.gradient_checkpointing and self.training: |
| 383 | layer_outputs = self._gradient_checkpointing_func( |
| 384 | layer_module.__call__, |
| 385 | hidden_states, |
| 386 | adaln_input, |
| 387 | layer_head_mask, |
| 388 | output_attentions, |
| 389 | ) |
| 390 | else: |
| 391 | layer_outputs = layer_module(hidden_states, adaln_input, layer_head_mask, output_attentions) |
| 392 | |
| 393 | hidden_states = layer_outputs[0] |
| 394 | |
| 395 | if output_attentions: |
| 396 | all_self_attentions = all_self_attentions + (layer_outputs[1],) |
| 397 | |
| 398 | if output_hidden_states: |
| 399 | all_hidden_states = all_hidden_states + (hidden_states,) |
| 400 | |
| 401 | if not return_dict: |
| 402 | return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) |
| 403 | return BaseModelOutput( |
| 404 | last_hidden_state=hidden_states, |
| 405 | hidden_states=all_hidden_states, |
| 406 | attentions=all_self_attentions, |
| 407 | ) |
| 408 | |
| 409 | |
| 410 | class ViTPreTrainedModel(PreTrainedModel): |