Transformer 2.0: Make metrics layer optional (#7075)

* trying fake merge call

* make metrics optional

* Remove extra print

official/transformer/v2/misc.py

@@ -102,6 +102,9 @@ def define_transformer_flags():
   flags.DEFINE_boolean(
       name='enable_tensorboard', default=False,
       help='Whether to enable Tensorboard callback.')
+  flags.DEFINE_boolean(
+      name='enable_metrics_in_training', default=False,
+      help='Whether to enable metrics during training.')
   flags.DEFINE_string(
       name='profile_steps', default=None,
       help='Save profiling data to model dir at given range of steps. The '

official/transformer/v2/transformer.py

@@ -42,7 +42,8 @@ def create_model(params, is_train):
       logits = internal_model([inputs, targets], training=is_train)
       vocab_size = params["vocab_size"]
       label_smoothing = params["label_smoothing"]
-      logits = metrics.MetricLayer(vocab_size)([logits, targets])
+      if params["enable_metrics_in_training"]:
+        logits = metrics.MetricLayer(vocab_size)([logits, targets])
       logits = metrics.LossLayer(vocab_size, label_smoothing)([logits, targets])
       logits = tf.keras.layers.Lambda(lambda x: x, name="logits")(logits)
       return tf.keras.Model([inputs, targets], logits)

official/transformer/v2/transformer_main.py

@@ -119,6 +119,7 @@ class TransformerTask(object):
     params["batch_size"] = flags_obj.batch_size or params["default_batch_size"]
     params["repeat_dataset"] = None
     params["dtype"] = flags_core.get_tf_dtype(flags_obj)
+    params["enable_metrics_in_training"] = flags_obj.enable_metrics_in_training
 
     if params["dtype"] == tf.float16:
       # TODO(reedwm): It's pretty ugly to set the global policy in a constructor