Allow user to specify root_data_dir in the benchmark class constructor (#6213)

* Allow user to specify root_data_dir in the benchmark class constructor

* Address comments

official/resnet/estimator_cifar_benchmark.py

@@ -29,22 +29,28 @@ import tensorflow as tf  # pylint: disable=g-bad-import-order
 from official.resnet import cifar10_main as cifar_main
 from official.utils.logs import hooks
 
-DATA_DIR = '/data/cifar10_data/cifar-10-batches-bin'
-
 
 class EstimatorCifar10BenchmarkTests(tf.test.Benchmark):
   """Benchmarks and accuracy tests for Estimator ResNet56."""
 
   local_flags = None
 
-  def __init__(self, output_dir=None):
+  def __init__(self, output_dir=None, root_data_dir=None):
+    """A benchmark class.
+
+    Args:
+      output_dir: directory where to output e.g. log files
+      root_data_dir: directory under which to look for dataset
+    """
+
     self.output_dir = output_dir
+    self.data_dir = os.path.join(root_data_dir, 'cifar-10-batches-bin')
 
   def resnet56_1_gpu(self):
     """Test layers model with Estimator and distribution strategies."""
     self._setup()
     flags.FLAGS.num_gpus = 1
-    flags.FLAGS.data_dir = DATA_DIR
+    flags.FLAGS.data_dir = self.data_dir
     flags.FLAGS.batch_size = 128
     flags.FLAGS.train_epochs = 182
     flags.FLAGS.model_dir = self._get_model_dir('resnet56_1_gpu')
@@ -57,7 +63,7 @@ class EstimatorCifar10BenchmarkTests(tf.test.Benchmark):
     """Test layers FP16 model with Estimator and distribution strategies."""
     self._setup()
     flags.FLAGS.num_gpus = 1
-    flags.FLAGS.data_dir = DATA_DIR
+    flags.FLAGS.data_dir = self.data_dir
     flags.FLAGS.batch_size = 128
     flags.FLAGS.train_epochs = 182
     flags.FLAGS.model_dir = self._get_model_dir('resnet56_fp16_1_gpu')
@@ -70,7 +76,7 @@ class EstimatorCifar10BenchmarkTests(tf.test.Benchmark):
     """Test layers model with Estimator and dist_strat. 2 GPUs."""
     self._setup()
     flags.FLAGS.num_gpus = 2
-    flags.FLAGS.data_dir = DATA_DIR
+    flags.FLAGS.data_dir = self.data_dir
     flags.FLAGS.batch_size = 128
     flags.FLAGS.train_epochs = 182
     flags.FLAGS.model_dir = self._get_model_dir('resnet56_2_gpu')
@@ -83,7 +89,7 @@ class EstimatorCifar10BenchmarkTests(tf.test.Benchmark):
     """Test layers FP16 model with Estimator and dist_strat. 2 GPUs."""
     self._setup()
     flags.FLAGS.num_gpus = 2
-    flags.FLAGS.data_dir = DATA_DIR
+    flags.FLAGS.data_dir = self.data_dir
     flags.FLAGS.batch_size = 128
     flags.FLAGS.train_epochs = 182
     flags.FLAGS.model_dir = self._get_model_dir('resnet56_fp16_2_gpu')
@@ -96,7 +102,7 @@ class EstimatorCifar10BenchmarkTests(tf.test.Benchmark):
     """A lightweight test that can finish quickly."""
     self._setup()
     flags.FLAGS.num_gpus = 1
-    flags.FLAGS.data_dir = DATA_DIR
+    flags.FLAGS.data_dir = self.data_dir
     flags.FLAGS.batch_size = 128
     flags.FLAGS.train_epochs = 1
     flags.FLAGS.model_dir = self._get_model_dir('resnet56_1_gpu')

official/resnet/keras/keras_cifar_benchmark.py

@@ -18,6 +18,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import os
 import time
 from absl import flags
 
@@ -26,7 +27,6 @@ from official.resnet.keras import keras_benchmark
 from official.resnet.keras import keras_cifar_main
 from official.resnet.keras import keras_common
 
-DATA_DIR = '/data/cifar10_data/cifar-10-batches-bin'
 MIN_TOP_1_ACCURACY = 0.925
 MAX_TOP_1_ACCURACY = 0.938
 
@@ -36,7 +36,15 @@ FLAGS = flags.FLAGS
 class Resnet56KerasAccuracy(keras_benchmark.KerasBenchmark):
   """Accuracy tests for ResNet56 Keras CIFAR-10."""
 
-  def __init__(self, output_dir=None):
+  def __init__(self, output_dir=None, root_data_dir=None):
+    """A benchmark class.
+
+    Args:
+      output_dir: directory where to output e.g. log files
+      root_data_dir: directory under which to look for dataset
+    """
+
+    self.data_dir = os.path.join(root_data_dir, 'cifar-10-batches-bin')
     flag_methods = [
         keras_common.define_keras_flags, cifar_main.define_cifar_flags
     ]
@@ -48,7 +56,7 @@ class Resnet56KerasAccuracy(keras_benchmark.KerasBenchmark):
     """Test keras based model with Keras fit and distribution strategies."""
     self._setup()
     FLAGS.num_gpus = 1
-    FLAGS.data_dir = DATA_DIR
+    FLAGS.data_dir = self.data_dir
     FLAGS.batch_size = 128
     FLAGS.train_epochs = 182
     FLAGS.model_dir = self._get_model_dir('benchmark_graph_1_gpu')
@@ -59,7 +67,7 @@ class Resnet56KerasAccuracy(keras_benchmark.KerasBenchmark):
     """Test keras based model with eager and distribution strategies."""
     self._setup()
     FLAGS.num_gpus = 1
-    FLAGS.data_dir = DATA_DIR
+    FLAGS.data_dir = self.data_dir
     FLAGS.batch_size = 128
     FLAGS.train_epochs = 182
     FLAGS.model_dir = self._get_model_dir('benchmark_1_gpu')
@@ -71,7 +79,7 @@ class Resnet56KerasAccuracy(keras_benchmark.KerasBenchmark):
     """Test keras based model with eager and distribution strategies."""
     self._setup()
     FLAGS.num_gpus = 2
-    FLAGS.data_dir = DATA_DIR
+    FLAGS.data_dir = self.data_dir
     FLAGS.batch_size = 128
     FLAGS.train_epochs = 182
     FLAGS.model_dir = self._get_model_dir('benchmark_2_gpu')
@@ -83,7 +91,7 @@ class Resnet56KerasAccuracy(keras_benchmark.KerasBenchmark):
     """Test keras based model with Keras fit and distribution strategies."""
     self._setup()
     FLAGS.num_gpus = 2
-    FLAGS.data_dir = DATA_DIR
+    FLAGS.data_dir = self.data_dir
     FLAGS.batch_size = 128
     FLAGS.train_epochs = 182
     FLAGS.model_dir = self._get_model_dir('benchmark_graph_2_gpu')
@@ -95,7 +103,7 @@ class Resnet56KerasAccuracy(keras_benchmark.KerasBenchmark):
     self._setup()
     FLAGS.distribution_strategy = 'off'
     FLAGS.num_gpus = 1
-    FLAGS.data_dir = DATA_DIR
+    FLAGS.data_dir = self.data_dir
     FLAGS.batch_size = 128
     FLAGS.train_epochs = 182
     FLAGS.model_dir = self._get_model_dir('benchmark_graph_1_gpu_no_dist_strat')
@@ -198,7 +206,7 @@ class Resnet56KerasBenchmarkBase(keras_benchmark.KerasBenchmark):
 class Resnet56KerasBenchmarkSynth(Resnet56KerasBenchmarkBase):
   """Synthetic benchmarks for ResNet56 and Keras."""
 
-  def __init__(self, output_dir=None):
+  def __init__(self, output_dir=None, root_data_dir=None):
     def_flags = {}
     def_flags['skip_eval'] = True
     def_flags['use_synthetic_data'] = True
@@ -212,10 +220,10 @@ class Resnet56KerasBenchmarkSynth(Resnet56KerasBenchmarkBase):
 class Resnet56KerasBenchmarkReal(Resnet56KerasBenchmarkBase):
   """Real data benchmarks for ResNet56 and Keras."""
 
-  def __init__(self, output_dir=None):
+  def __init__(self, output_dir=None, root_data_dir=None):
     def_flags = {}
     def_flags['skip_eval'] = True
-    def_flags['data_dir'] = DATA_DIR
+    def_flags['data_dir'] = self.data_dir
     def_flags['train_steps'] = 110
     def_flags['log_steps'] = 10
 

official/resnet/keras/keras_imagenet_benchmark.py

@@ -27,7 +27,6 @@ from official.resnet.keras import keras_imagenet_main
 
 MIN_TOP_1_ACCURACY = 0.76
 MAX_TOP_1_ACCURACY = 0.77
-DATA_DIR = '/data/imagenet/'
 
 FLAGS = flags.FLAGS
 
@@ -35,11 +34,19 @@ FLAGS = flags.FLAGS
 class Resnet50KerasAccuracy(keras_benchmark.KerasBenchmark):
   """Benchmark accuracy tests for ResNet50 in Keras."""
 
-  def __init__(self, output_dir=None):
+  def __init__(self, output_dir=None, root_data_dir=None):
+    """A benchmark class.
+
+    Args:
+      output_dir: directory where to output e.g. log files
+      root_data_dir: directory under which to look for dataset
+    """
+
     flag_methods = [
         keras_common.define_keras_flags, imagenet_main.define_imagenet_flags
     ]
 
+    self.data_dir = os.path.join(root_data_dir, 'imagenet')
     super(Resnet50KerasAccuracy, self).__init__(
         output_dir=output_dir, flag_methods=flag_methods)
 
@@ -47,7 +54,7 @@ class Resnet50KerasAccuracy(keras_benchmark.KerasBenchmark):
     """Test Keras model with Keras fit/dist_strat and 8 GPUs."""
     self._setup()
     FLAGS.num_gpus = 8
-    FLAGS.data_dir = DATA_DIR
+    FLAGS.data_dir = self.data_dir
     FLAGS.batch_size = 128 * 8
     FLAGS.train_epochs = 90
     FLAGS.model_dir = self._get_model_dir('benchmark_graph_8_gpu')
@@ -58,7 +65,7 @@ class Resnet50KerasAccuracy(keras_benchmark.KerasBenchmark):
     """Test Keras model with eager, dist_strat and 8 GPUs."""
     self._setup()
     FLAGS.num_gpus = 8
-    FLAGS.data_dir = DATA_DIR
+    FLAGS.data_dir = self.data_dir
     FLAGS.batch_size = 128 * 8
     FLAGS.train_epochs = 90
     FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu')
@@ -70,7 +77,7 @@ class Resnet50KerasAccuracy(keras_benchmark.KerasBenchmark):
     """Restricts CPU memory allocation."""
     self._setup()
     FLAGS.num_gpus = 8
-    FLAGS.data_dir = DATA_DIR
+    FLAGS.data_dir = self.data_dir
     FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu_bfc_allocator')
     FLAGS.dtype = 'fp32'
     FLAGS.batch_size = 128 * 8  # 8 GPUs
@@ -211,7 +218,7 @@ class Resnet50KerasBenchmarkReal(Resnet50KerasBenchmarkBase):
   def __init__(self, output_dir=None):
     def_flags = {}
     def_flags['skip_eval'] = True
-    def_flags['data_dir'] = DATA_DIR
+    def_flags['data_dir'] = self.data_dir
     def_flags['train_steps'] = 110
     def_flags['log_steps'] = 10