| 278 | |
| 279 | |
| 280 | class BertAttention(nn.Module): |
| 281 | def __init__(self, config): |
| 282 | super().__init__() |
| 283 | self.self = BertSelfAttention(config) |
| 284 | self.output = BertSelfOutput(config) |
| 285 | self.pruned_heads = set() |
| 286 | |
| 287 | def prune_heads(self, heads): |
| 288 | if len(heads) == 0: |
| 289 | return |
| 290 | heads, index = find_pruneable_heads_and_indices( |
| 291 | heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads |
| 292 | ) |
| 293 | |
| 294 | # Prune linear layers |
| 295 | self.self.query = prune_linear_layer(self.self.query, index) |
| 296 | self.self.key = prune_linear_layer(self.self.key, index) |
| 297 | self.self.value = prune_linear_layer(self.self.value, index) |
| 298 | self.output.dense = prune_linear_layer(self.output.dense, index, dim=1) |
| 299 | |
| 300 | # Update hyper params and store pruned heads |
| 301 | self.self.num_attention_heads = self.self.num_attention_heads - len(heads) |
| 302 | self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads |
| 303 | self.pruned_heads = self.pruned_heads.union(heads) |
| 304 | |
| 305 | def forward( |
| 306 | self, |
| 307 | hidden_states, |
| 308 | attention_mask=None, |
| 309 | head_mask=None, |
| 310 | encoder_hidden_states=None, |
| 311 | encoder_attention_mask=None, |
| 312 | output_attentions=False, |
| 313 | ): |
| 314 | self_outputs = self.self( |
| 315 | hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, output_attentions, |
| 316 | ) |
| 317 | attention_output = self.output(self_outputs[0], hidden_states) |
| 318 | outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them |
| 319 | return outputs |
| 320 | |
| 321 | |
| 322 | class BertIntermediate(nn.Module): |