SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

prune both attention and self-attention heads

Browse Source
This commit is contained in:
Rémi Louf
2019-10-10 12:17:22 +02:00
parent 17177e7379
commit 51261167b4
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
transformers/modeling_bert.py
View File

@@ -633,7 +633,7 @@ class BertModel(BertPreTrainedModel):
See base class PreTrainedModel
"""
for layer, heads in heads_to_prune.items():
self.encoder.layer[layer].attention.prune_heads(heads)
self.encoder.layer[layer].self_attention.prune_heads(heads)
def forward(self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None):
if attention_mask is None:
@@ -736,7 +736,8 @@ class BertDecoderModel(BertPreTrainedModel):
See base class PreTrainedModel
"""
for layer, heads in heads_to_prune.items():
self.encoder.layer[layer].attention.prune_heads(heads)
self.decoder.layer[layer].attention.prune_heads(heads)
self.decoder.layer[layer].self_attention.prune_heads(heads)
def forward(self, input_ids, encoder_outputs, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None):
if attention_mask is None: