Internal change

PiperOrigin-RevId: 331906632

official/nlp/tasks/masked_lm.py

@@ -63,31 +63,33 @@ class MaskedLMTask(base_task.Task):
                    model_outputs,
                    metrics,
                    aux_losses=None) -> tf.Tensor:
-    metrics = dict([(metric.name, metric) for metric in metrics])
-    lm_prediction_losses = tf.keras.losses.sparse_categorical_crossentropy(
-        labels['masked_lm_ids'],
-        tf.cast(model_outputs['lm_output'], tf.float32),
-        from_logits=True)
-    lm_label_weights = labels['masked_lm_weights']
-    lm_numerator_loss = tf.reduce_sum(lm_prediction_losses * lm_label_weights)
-    lm_denominator_loss = tf.reduce_sum(lm_label_weights)
-    mlm_loss = tf.math.divide_no_nan(lm_numerator_loss, lm_denominator_loss)
-    metrics['lm_example_loss'].update_state(mlm_loss)
-    if 'next_sentence_labels' in labels:
-      sentence_labels = labels['next_sentence_labels']
-      sentence_outputs = tf.cast(
-          model_outputs['next_sentence'], dtype=tf.float32)
-      sentence_loss = tf.reduce_mean(
-          tf.keras.losses.sparse_categorical_crossentropy(
-              sentence_labels, sentence_outputs, from_logits=True))
-      metrics['next_sentence_loss'].update_state(sentence_loss)
-      total_loss = mlm_loss + sentence_loss
-    else:
-      total_loss = mlm_loss
-
-    if aux_losses:
-      total_loss += tf.add_n(aux_losses)
-    return total_loss
+    with tf.name_scope('MaskedLMTask/losses'):
+      metrics = dict([(metric.name, metric) for metric in metrics])
+      lm_prediction_losses = tf.keras.losses.sparse_categorical_crossentropy(
+          labels['masked_lm_ids'],
+          tf.cast(model_outputs['lm_output'], tf.float32),
+          from_logits=True)
+      lm_label_weights = labels['masked_lm_weights']
+      lm_numerator_loss = tf.reduce_sum(lm_prediction_losses *
+                                        lm_label_weights)
+      lm_denominator_loss = tf.reduce_sum(lm_label_weights)
+      mlm_loss = tf.math.divide_no_nan(lm_numerator_loss, lm_denominator_loss)
+      metrics['lm_example_loss'].update_state(mlm_loss)
+      if 'next_sentence_labels' in labels:
+        sentence_labels = labels['next_sentence_labels']
+        sentence_outputs = tf.cast(
+            model_outputs['next_sentence'], dtype=tf.float32)
+        sentence_loss = tf.reduce_mean(
+            tf.keras.losses.sparse_categorical_crossentropy(
+                sentence_labels, sentence_outputs, from_logits=True))
+        metrics['next_sentence_loss'].update_state(sentence_loss)
+        total_loss = mlm_loss + sentence_loss
+      else:
+        total_loss = mlm_loss
+
+      if aux_losses:
+        total_loss += tf.add_n(aux_losses)
+      return total_loss
 
   def build_inputs(self, params, input_context=None):
     """Returns tf.data.Dataset for pretraining."""
@@ -128,14 +130,15 @@ class MaskedLMTask(base_task.Task):
     return metrics
 
   def process_metrics(self, metrics, labels, model_outputs):
-    metrics = dict([(metric.name, metric) for metric in metrics])
-    if 'masked_lm_accuracy' in metrics:
-      metrics['masked_lm_accuracy'].update_state(labels['masked_lm_ids'],
-                                                 model_outputs['lm_output'],
-                                                 labels['masked_lm_weights'])
-    if 'next_sentence_accuracy' in metrics:
-      metrics['next_sentence_accuracy'].update_state(
-          labels['next_sentence_labels'], model_outputs['next_sentence'])
+    with tf.name_scope('MaskedLMTask/process_metrics'):
+      metrics = dict([(metric.name, metric) for metric in metrics])
+      if 'masked_lm_accuracy' in metrics:
+        metrics['masked_lm_accuracy'].update_state(
+            labels['masked_lm_ids'], model_outputs['lm_output'],
+            labels['masked_lm_weights'])
+      if 'next_sentence_accuracy' in metrics:
+        metrics['next_sentence_accuracy'].update_state(
+            labels['next_sentence_labels'], model_outputs['next_sentence'])
 
   def train_step(self, inputs, model: tf.keras.Model,
                  optimizer: tf.keras.optimizers.Optimizer, metrics):