Internal change

PiperOrigin-RevId: 308062306

official/benchmark/perfzero_benchmark.py

@@ -39,6 +39,7 @@ class PerfZeroBenchmark(tf.test.Benchmark):
   def __init__(self,
                output_dir=None,
                default_flags=None,
+               root_data_dir=None,
                flag_methods=None,
                tpu=None):
     """Initialize class.
@@ -46,6 +47,8 @@ class PerfZeroBenchmark(tf.test.Benchmark):
     Args:
       output_dir: Base directory to store all output for the test.
       default_flags: Set of flags to pass to model.
+      root_data_dir: Optional param used by child classes to look for the
+        dataset.
       flag_methods: Set of flag methods to run during setup.
       tpu: (optional) TPU name to use in a TPU benchmark.
     """
@@ -70,6 +73,8 @@ class PerfZeroBenchmark(tf.test.Benchmark):
       # the TPU at runtime and passes the TPU's name to this flag.
       self.default_flags['tpu'] = resolved_tpu
 
+    logging.info('root_data_dir: %s', root_data_dir)
+
   def _get_model_dir(self, folder_name):
     """Returns directory to store info, e.g. saved model and event log."""
     return os.path.join(self.output_dir, folder_name)