Disabling Tensorboard profiling for NCF. (#9959)

PiperOrigin-RevId: 371256980

Co-authored-by: A. Unique TensorFlower <[email protected]>

Author: gagika

official/recommendation/ncf_keras_main.py

@@ -258,10 +258,9 @@ def run_ncf(_):
         "val_HR_METRIC", desired_value=FLAGS.hr_threshold)
     callbacks.append(early_stopping_callback)
 
-  (train_input_dataset, eval_input_dataset,
-   num_train_steps, num_eval_steps) = \
-    (ncf_input_pipeline.create_ncf_input_data(
-        params, producer, input_meta_data, strategy))
+  (train_input_dataset, eval_input_dataset, num_train_steps,
+   num_eval_steps) = ncf_input_pipeline.create_ncf_input_data(
+       params, producer, input_meta_data, strategy)
   steps_per_epoch = None if generate_input_online else num_train_steps
 
   with distribute_utils.get_strategy_scope(strategy):
@@ -307,7 +306,8 @@ def run_ncf(_):
       if not FLAGS.ml_perf:
         # Create Tensorboard summary and checkpoint callbacks.
         summary_dir = os.path.join(FLAGS.model_dir, "summaries")
-        summary_callback = tf.keras.callbacks.TensorBoard(summary_dir)
+        summary_callback = tf.keras.callbacks.TensorBoard(
+            summary_dir, profile_batch=0)
         checkpoint_path = os.path.join(FLAGS.model_dir, "checkpoint")
         checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
             checkpoint_path, save_weights_only=True)

official/recommendation/ncf_test.py

@@ -38,11 +38,13 @@ def setUpClass(cls):  # pylint: disable=invalid-name
     ncf_common.define_ncf_flags()
 
   def setUp(self):
+    super().setUp()
     self.top_k_old = rconst.TOP_K
     self.num_eval_negatives_old = rconst.NUM_EVAL_NEGATIVES
     rconst.NUM_EVAL_NEGATIVES = 2
 
   def tearDown(self):
+    super().tearDown()
     rconst.NUM_EVAL_NEGATIVES = self.num_eval_negatives_old
     rconst.TOP_K = self.top_k_old