| 459 | |
| 460 | |
| 461 | class BertEncoder(nn.Module): |
| 462 | def __init__(self, config): |
| 463 | super().__init__() |
| 464 | self.config = config |
| 465 | self.layer = nn.ModuleList([BertLayer(config,i) for i in range(config.num_hidden_layers)]) |
| 466 | self.gradient_checkpointing = False |
| 467 | |
| 468 | def forward( |
| 469 | self, |
| 470 | hidden_states, |
| 471 | attention_mask=None, |
| 472 | head_mask=None, |
| 473 | encoder_hidden_states=None, |
| 474 | encoder_attention_mask=None, |
| 475 | past_key_values=None, |
| 476 | use_cache=None, |
| 477 | output_attentions=False, |
| 478 | output_hidden_states=False, |
| 479 | return_dict=True, |
| 480 | mode='multimodal', |
| 481 | ): |
| 482 | all_hidden_states = () if output_hidden_states else None |
| 483 | all_self_attentions = () if output_attentions else None |
| 484 | all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None |
| 485 | |
| 486 | next_decoder_cache = () if use_cache else None |
| 487 | |
| 488 | for i in range(self.config.num_hidden_layers): |
| 489 | layer_module = self.layer[i] |
| 490 | if output_hidden_states: |
| 491 | all_hidden_states = all_hidden_states + (hidden_states,) |
| 492 | |
| 493 | layer_head_mask = head_mask[i] if head_mask is not None else None |
| 494 | past_key_value = past_key_values[i] if past_key_values is not None else None |
| 495 | |
| 496 | if self.gradient_checkpointing and self.training: |
| 497 | |
| 498 | if use_cache: |
| 499 | logger.warn( |
| 500 | "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." |
| 501 | ) |
| 502 | use_cache = False |
| 503 | |
| 504 | def create_custom_forward(module): |
| 505 | def custom_forward(*inputs): |
| 506 | return module(*inputs, past_key_value, output_attentions) |
| 507 | |
| 508 | return custom_forward |
| 509 | |
| 510 | layer_outputs = torch.utils.checkpoint.checkpoint( |
| 511 | create_custom_forward(layer_module), |
| 512 | hidden_states, |
| 513 | attention_mask, |
| 514 | layer_head_mask, |
| 515 | encoder_hidden_states, |
| 516 | encoder_attention_mask, |
| 517 | mode=mode, |
| 518 | ) |