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

Head pruning for ALBERT

Browse Source
This commit is contained in:
Lysandre
2019-11-07 17:09:16 +00:00
committed by Lysandre Debut
parent 4374eaea78
commit abb23a78ba
2 changed files with 49 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
transformers/tests/modeling_albert_test.py
View File

@@ -35,7 +35,6 @@ else:
class AlbertModelTest(CommonTestCases.CommonModelTester):
all_model_classes = (AlbertModel, AlbertForMaskedLM) if is_torch_available() else ()
test_pruning = False
test_head_masking = False
class AlbertModelTester(object):
@@ -49,9 +48,10 @@ class AlbertModelTest(CommonTestCases.CommonModelTester):
use_token_type_ids=True,
use_labels=True,
vocab_size=99,
hidden_size=32,
num_hidden_layers=5,
num_attention_heads=4,
hidden_size=36,
num_hidden_layers=6,
num_hidden_groups=6,
num_attention_heads=6,
intermediate_size=37,
hidden_act="gelu",
hidden_dropout_prob=0.1,
@@ -86,6 +86,7 @@ class AlbertModelTest(CommonTestCases.CommonModelTester):
self.num_labels = num_labels
self.num_choices = num_choices
self.scope = scope
self.num_hidden_groups = num_hidden_groups
def prepare_config_and_inputs(self):
input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
@@ -117,7 +118,8 @@ class AlbertModelTest(CommonTestCases.CommonModelTester):
attention_probs_dropout_prob=self.attention_probs_dropout_prob,
max_position_embeddings=self.max_position_embeddings,
type_vocab_size=self.type_vocab_size,
initializer_range=self.initializer_range)
initializer_range=self.initializer_range,
num_hidden_groups=self.num_hidden_groups)
return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels