From 5d80539488705ae1de95b93ed86dfd1414075331 Mon Sep 17 00:00:00 2001 From: Julien Plu Date: Thu, 12 Nov 2020 20:08:26 +0100 Subject: [PATCH] Add pretraining loss computation for TF Bert pretraining (#8470) * Add pretraining loss computation for TF Bert pretraining * Fix labels creation * Fix T5 model * restore T5 kwargs * try a generic fix for pretraining models * Apply style * Overide the prepare method for the BERT tests --- src/transformers/modeling_tf_bert.py | 84 ++++++++++++++++++++++++++-- tests/test_modeling_tf_bert.py | 11 ++++ tests/test_modeling_tf_common.py | 27 ++++++--- 3 files changed, 109 insertions(+), 13 deletions(-) diff --git a/src/transformers/modeling_tf_bert.py b/src/transformers/modeling_tf_bert.py index baf1d962c0..72da644ebc 100644 --- a/src/transformers/modeling_tf_bert.py +++ b/src/transformers/modeling_tf_bert.py @@ -89,6 +89,38 @@ TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST = [ ] +class TFBertPreTrainingLoss: + """ + Loss function suitable for BERT-like pre-training, that is, the task of pretraining a language model by combining + NSP + MLM. .. note:: Any label of -100 will be ignored (along with the corresponding logits) in the loss + computation. + """ + + def compute_loss(self, labels, logits): + loss_fn = tf.keras.losses.SparseCategoricalCrossentropy( + from_logits=True, reduction=tf.keras.losses.Reduction.NONE + ) + # make sure only labels that are not equal to -100 + # are taken into account as loss + masked_lm_active_loss = tf.not_equal(tf.reshape(labels["labels"], (-1,)), -100) + masked_lm_reduced_logits = tf.boolean_mask( + tf.reshape(logits[0], (-1, shape_list(logits[0])[2])), + masked_lm_active_loss, + ) + masked_lm_labels = tf.boolean_mask(tf.reshape(labels["labels"], (-1,)), masked_lm_active_loss) + next_sentence_active_loss = tf.not_equal(tf.reshape(labels["next_sentence_label"], (-1,)), -100) + next_sentence_reduced_logits = tf.boolean_mask(tf.reshape(logits[1], (-1, 2)), next_sentence_active_loss) + next_sentence_label = tf.boolean_mask( + tf.reshape(labels["next_sentence_label"], (-1,)), mask=next_sentence_active_loss + ) + masked_lm_loss = loss_fn(masked_lm_labels, masked_lm_reduced_logits) + next_sentence_loss = loss_fn(next_sentence_label, next_sentence_reduced_logits) + masked_lm_loss = tf.reshape(masked_lm_loss, (-1, shape_list(next_sentence_loss)[0])) + masked_lm_loss = tf.reduce_mean(masked_lm_loss, 0) + + return masked_lm_loss + next_sentence_loss + + class TFBertEmbeddings(tf.keras.layers.Layer): """Construct the embeddings from word, position and token_type embeddings.""" @@ -688,6 +720,7 @@ class TFBertForPreTrainingOutput(ModelOutput): heads. """ + loss: Optional[tf.Tensor] = None prediction_logits: tf.Tensor = None seq_relationship_logits: tf.Tensor = None hidden_states: Optional[Tuple[tf.Tensor]] = None @@ -814,7 +847,7 @@ Bert Model with two heads on top as done during the pre-training: """, BERT_START_DOCSTRING, ) -class TFBertForPreTraining(TFBertPreTrainedModel): +class TFBertForPreTraining(TFBertPreTrainedModel, TFBertPreTrainingLoss): def __init__(self, config, *inputs, **kwargs): super().__init__(config, *inputs, **kwargs) @@ -827,7 +860,21 @@ class TFBertForPreTraining(TFBertPreTrainedModel): @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length")) @replace_return_docstrings(output_type=TFBertForPreTrainingOutput, config_class=_CONFIG_FOR_DOC) - def call(self, inputs, **kwargs): + def call( + self, + inputs=None, + attention_mask=None, + token_type_ids=None, + position_ids=None, + head_mask=None, + inputs_embeds=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + labels=None, + next_sentence_label=None, + training=False, + ): r""" Return: @@ -843,17 +890,44 @@ class TFBertForPreTraining(TFBertPreTrainedModel): >>> prediction_scores, seq_relationship_scores = outputs[:2] """ - return_dict = kwargs.get("return_dict") return_dict = return_dict if return_dict is not None else self.bert.return_dict - outputs = self.bert(inputs, **kwargs) + + if isinstance(inputs, (tuple, list)): + labels = inputs[9] if len(inputs) > 9 else labels + next_sentence_label = inputs[10] if len(inputs) > 10 else next_sentence_label + if len(inputs) > 9: + inputs = inputs[:9] + elif isinstance(inputs, (dict, BatchEncoding)): + labels = inputs.pop("labels", labels) + next_sentence_label = inputs.pop("next_sentence_label", next_sentence_label) + + outputs = self.bert( + inputs, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + training=training, + ) sequence_output, pooled_output = outputs[:2] - prediction_scores = self.mlm(sequence_output, training=kwargs.get("training", False)) + prediction_scores = self.mlm(sequence_output, training=training) seq_relationship_score = self.nsp(pooled_output) + total_loss = None + + if labels is not None and next_sentence_label is not None: + d_labels = {"labels": labels} + d_labels["next_sentence_label"] = next_sentence_label + total_loss = self.compute_loss(labels=d_labels, logits=(prediction_scores, seq_relationship_score)) if not return_dict: return (prediction_scores, seq_relationship_score) + outputs[2:] return TFBertForPreTrainingOutput( + loss=total_loss, prediction_logits=prediction_scores, seq_relationship_logits=seq_relationship_score, hidden_states=outputs.hidden_states, diff --git a/tests/test_modeling_tf_bert.py b/tests/test_modeling_tf_bert.py index 0541637499..48a14f4a23 100644 --- a/tests/test_modeling_tf_bert.py +++ b/tests/test_modeling_tf_bert.py @@ -26,6 +26,7 @@ from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor if is_tf_available(): import tensorflow as tf + from transformers import TF_MODEL_FOR_PRETRAINING_MAPPING from transformers.modeling_tf_bert import ( TFBertForMaskedLM, TFBertForMultipleChoice, @@ -274,6 +275,16 @@ class TFBertModelTest(TFModelTesterMixin, unittest.TestCase): else () ) + # special case for ForPreTraining model + def _prepare_for_class(self, inputs_dict, model_class, return_labels=False): + inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels) + + if return_labels: + if model_class in TF_MODEL_FOR_PRETRAINING_MAPPING.values(): + inputs_dict["next_sentence_label"] = tf.zeros(self.model_tester.batch_size, dtype=tf.int32) + + return inputs_dict + def setUp(self): self.model_tester = TFBertModelTester(self) self.config_tester = ConfigTester(self, config_class=BertConfig, hidden_size=37) diff --git a/tests/test_modeling_tf_common.py b/tests/test_modeling_tf_common.py index 30321fc5e6..c89acc3983 100644 --- a/tests/test_modeling_tf_common.py +++ b/tests/test_modeling_tf_common.py @@ -36,6 +36,7 @@ if is_tf_available(): TF_MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING, TF_MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING, + TF_MODEL_FOR_PRETRAINING_MAPPING, TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING, TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING, @@ -102,6 +103,7 @@ class TFModelTesterMixin: *TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.values(), *TF_MODEL_FOR_CAUSAL_LM_MAPPING.values(), *TF_MODEL_FOR_MASKED_LM_MAPPING.values(), + *TF_MODEL_FOR_PRETRAINING_MAPPING.values(), *TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.values(), ]: inputs_dict["labels"] = tf.zeros( @@ -834,7 +836,9 @@ class TFModelTesterMixin: if getattr(model, "compute_loss", None): # The number of elements in the loss should be the same as the number of elements in the label prepared_for_class = self._prepare_for_class(inputs_dict.copy(), model_class, return_labels=True) - added_label = prepared_for_class[list(prepared_for_class.keys() - inputs_dict.keys())[0]] + added_label = prepared_for_class[ + sorted(list(prepared_for_class.keys() - inputs_dict.keys()), reverse=True)[0] + ] loss_size = tf.size(added_label) if model.__class__ in TF_MODEL_FOR_CAUSAL_LM_MAPPING.values(): @@ -859,23 +863,30 @@ class TFModelTesterMixin: # Get keys that were added with the _prepare_for_class function label_keys = prepared_for_class.keys() - inputs_dict.keys() - signature = inspect.getfullargspec(model.call)[0] + signature = inspect.signature(model.call).parameters + signature_names = list(signature.keys()) # Create a dictionary holding the location of the tensors in the tuple - tuple_index_mapping = {1: "input_ids"} + tuple_index_mapping = {0: "input_ids"} for label_key in label_keys: - label_key_index = signature.index(label_key) + label_key_index = signature_names.index(label_key) tuple_index_mapping[label_key_index] = label_key sorted_tuple_index_mapping = sorted(tuple_index_mapping.items()) + # Initialize a list with their default values, update the values and convert to a tuple + list_input = [] + + for name in signature_names: + if name != "kwargs": + list_input.append(signature[name].default) - # Initialize a list with None, update the values and convert to a tuple - list_input = [None] * sorted_tuple_index_mapping[-1][0] for index, value in sorted_tuple_index_mapping: - list_input[index - 1] = prepared_for_class[value] + list_input[index] = prepared_for_class[value] + tuple_input = tuple(list_input) # Send to model - loss = model(tuple_input)[0] + loss = model(tuple_input[:-1])[0] + self.assertEqual(loss.shape, [loss_size]) def _generate_random_bad_tokens(self, num_bad_tokens, model):