Add benchmark tests and unit tests for clustering mobilenet_v1

official/benchmark/keras_imagenet_benchmark.py

@@ -42,6 +42,7 @@ MOBILENET_V1_MAX_TOP_1_ACCURACY = 0.68
 MODEL_OPTIMIZATION_TOP_1_ACCURACY = {
     'RESNET50_FINETUNE_PRUNING': (0.76, 0.77),
     'MOBILENET_V1_FINETUNE_PRUNING': (0.67, 0.68),
+    'MOBILENET_V1_FINETUNE_CLUSTERING': (0.68, 0.70)
 }
 
 FLAGS = flags.FLAGS
@@ -1774,5 +1775,119 @@ class Resnet50KerasPruningBenchmarkReal(KerasPruningBenchmarkRealBase):
         default_flags=default_flags, **kwargs)
 
 
+class KerasClusteringAccuracyBase(keras_benchmark.KerasBenchmark):
+  """Benchmark accuracy tests for clustering method."""
+
+  def __init__(self,
+               output_dir=None,
+               root_data_dir=None,
+               default_flags=None,
+               **kwargs):
+    """An accuracy benchmark class for clustering method.
+
+    Args:
+      output_dir: directory where to output e.g. log files
+      root_data_dir: directory under which to look for dataset
+      default_flags: default flags
+      **kwargs: arbitrary named arguments. This is needed to make the
+                constructor forward compatible in case PerfZero provides more
+                named arguments before updating the constructor.
+    """
+    if default_flags is None:
+      default_flags = {}
+    default_flags['clustering_method'] = 'selective_clustering'
+    default_flags['data_dir'] = os.path.join(root_data_dir, 'imagenet')
+    default_flags['model'] = 'mobilenet_pretrained'
+    default_flags['optimizer'] = 'mobilenet_fine_tune'
+
+    flag_methods = [resnet_imagenet_main.define_imagenet_keras_flags]
+
+    super(KerasClusteringAccuracyBase, self).__init__(
+        output_dir=output_dir,
+        flag_methods=flag_methods,
+        default_flags=default_flags,
+        **kwargs)
+
+  def benchmark_8_gpu(self):
+    """Test Keras model with eager, dist_strat and 8 GPUs."""
+    self._setup()
+
+    FLAGS.num_gpus = 8
+    FLAGS.batch_size = 32 * 8
+    FLAGS.train_epochs = 1
+    FLAGS.model_dir = self._get_model_dir('benchmark_8_gpu')
+    FLAGS.dtype = 'fp32'
+    FLAGS.enable_eager = True
+    self._run_and_report_benchmark()
+
+  @benchmark_wrappers.enable_runtime_flags
+  def _run_and_report_benchmark(self,
+                                top_1_min=MODEL_OPTIMIZATION_TOP_1_ACCURACY[
+                                    'MOBILENET_V1_FINETUNE_CLUSTERING'][0],
+                                top_1_max=MODEL_OPTIMIZATION_TOP_1_ACCURACY[
+                                    'MOBILENET_V1_FINETUNE_CLUSTERING'][1]):
+    start_time_sec = time.time()
+    stats = resnet_imagenet_main.run(flags.FLAGS)
+    wall_time_sec = time.time() - start_time_sec
+
+    super(KerasClusteringAccuracyBase, self)._report_benchmark(
+        stats,
+        wall_time_sec,
+        top_1_min=top_1_min,
+        top_1_max=top_1_max,
+        total_batch_size=FLAGS.batch_size,
+        log_steps=100)
+
+
+class MobilenetV1KerasClusteringAccuracy(KerasClusteringAccuracyBase):
+  """Benchmark accuracy tests for MobilenetV1 with clustering method."""
+
+  def __init__(self, root_data_dir=None, **kwargs):
+    default_flags = {
+        'model': 'mobilenet_pretrained',
+        'optimizer': 'mobilenet_fine_tune',
+    }
+    super(MobilenetV1KerasClusteringAccuracy, self).__init__(
+        root_data_dir=root_data_dir,
+        default_flags=default_flags,
+        **kwargs)
+
+  def _run_and_report_benchmark(self):
+    super(MobilenetV1KerasClusteringAccuracy, self)._run_and_report_benchmark(
+        top_1_min=\
+        MODEL_OPTIMIZATION_TOP_1_ACCURACY['MOBILENET_V1_FINETUNE_CLUSTERING'][0],
+        top_1_max=\
+        MODEL_OPTIMIZATION_TOP_1_ACCURACY['MOBILENET_V1_FINETUNE_CLUSTERING'][1])
+
+
+class KerasClusteringBenchmarkRealBase(Resnet50KerasBenchmarkBase):
+  """Clustering method benchmarks."""
+
+  def __init__(self, root_data_dir=None, default_flags=None, **kwargs):
+    if default_flags is None:
+      default_flags = {}
+    default_flags.update({
+        'skip_eval': True,
+        'report_accuracy_metrics': False,
+        'data_dir': os.path.join(root_data_dir, 'imagenet'),
+        'train_steps': 110,
+        'log_steps': 10,
+    })
+    super(KerasClusteringBenchmarkRealBase, self).__init__(
+        default_flags=default_flags, **kwargs)
+
+
+class MobilenetV1KerasClusteringBenchmarkReal(KerasClusteringBenchmarkRealBase):
+  """Clustering method benchmarks for MobilenetV1."""
+
+  def __init__(self, **kwargs):
+    default_flags = {
+        'model': 'mobilenet_pretrained',
+        'optimizer': 'mobilenet_fine_tune',
+    }
+    super(MobilenetV1KerasClusteringBenchmarkReal, self).__init__(
+        default_flags=default_flags, **kwargs)
+
+
 if __name__ == '__main__':
   tf.test.main()

official/benchmark/models/resnet_imagenet_main.py

@@ -63,7 +63,7 @@ def selective_clustering_clone_wrapper(clustering_params1, clustering_params2,
         return cluster.cluster_weights(layer, **clustering_params1)
       else:
         print("Wrapped layer " + layer.name +
-              " with number of clusters equals to " +
+              " with " +
               str(clustering_params2["number_of_clusters"]) + " clusters.")
         return cluster.cluster_weights(layer, **clustering_params2)
     return layer
@@ -243,7 +243,7 @@ def run(flags_obj):
           classes=imagenet_preprocessing.NUM_CLASSES,
           layers=tf.keras.layers)
     elif flags_obj.model == 'mobilenet_pretrained':
-      shape = (3, 224, 224)
+      shape = (224, 224, 3)
       model = tf.keras.applications.mobilenet.MobileNet(
           input_shape=shape,
           alpha=1.0,
@@ -366,12 +366,13 @@ def run(flags_obj):
   if not strategy and flags_obj.explicit_gpu_placement:
     no_dist_strat_device.__exit__()
 
-  if flags_obj.save_files_to:
-    keras_file = os.path.join(flags_obj.save_files_to, 'clustered.h5')
-  else:
-    keras_file = './clustered.h5'
-  print('Saving clustered and stripped model to: ', keras_file)
-  tf.keras.models.save_model(model, keras_file)
+  if flags_obj.clustering_method:
+    if flags_obj.save_files_to:
+      keras_file = os.path.join(flags_obj.save_files_to, 'clustered.h5')
+    else:
+      keras_file = './clustered.h5'
+    print('Saving clustered and stripped model to: ', keras_file)
+    tf.keras.models.save_model(model, keras_file)
 
   stats = common.build_stats(history, eval_output, callbacks)
   return stats

official/benchmark/models/resnet_imagenet_test.py

@@ -26,12 +26,13 @@ from official.benchmark.models import resnet_imagenet_main
 from official.utils.testing import integration
 from official.vision.image_classification.resnet import imagenet_preprocessing
 
-
+# TBC: joint clustering and tuning is not supported yet so only one flag should be selected
 @parameterized.parameters(
     "resnet",
     # "resnet_polynomial_decay",  b/151854314
     "mobilenet",
-    # "mobilenet_polynomial_decay"  b/151854314
+    # "mobilenet_polynomial_decay",  b/151854314
+    "mobilenet_selective_clustering",
 )
 class KerasImagenetTest(tf.test.TestCase):
   """Unit tests for Keras Models with ImageNet."""
@@ -74,6 +75,11 @@ class KerasImagenetTest(tf.test.TestCase):
           "-pruning_method",
           "polynomial_decay",
       ],
+      "mobilenet_selective_clustering": [
+          "-model", "mobilenet_pretrained",
+          "-optimizer", "mobilenet_fine_tune",
+          "-clustering_method", "selective_clustering",
+      ]
   }
   _tempdir = None
 
@@ -167,7 +173,10 @@ class KerasImagenetTest(tf.test.TestCase):
     extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)
 
     if "polynomial_decay" in extra_flags:
-      self.skipTest("Pruning with fp16 is not currently supported.")
+      self.skipTest("Pruning with fp16 is currently not supported.")
+
+    if "selective_clustering" in extra_flags:
+      self.skipTest("Clustering with fp16 is currently not supported.")
 
     integration.run_synthetic(
         main=resnet_imagenet_main.run,
@@ -237,7 +246,10 @@ class KerasImagenetTest(tf.test.TestCase):
     extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)
 
     if "polynomial_decay" in extra_flags:
-      self.skipTest("Pruning with fp16 is not currently supported.")
+      self.skipTest("Pruning with fp16 is currently not supported.")
+
+    if "selective_clustering" in extra_flags:
+      self.skipTest("Clustering with fp16 is currently not supported.")
 
     integration.run_synthetic(
         main=resnet_imagenet_main.run,
@@ -264,7 +276,10 @@ class KerasImagenetTest(tf.test.TestCase):
     extra_flags = extra_flags + self.get_extra_flags_dict(flags_key)
 
     if "polynomial_decay" in extra_flags:
-      self.skipTest("Pruning with fp16 is not currently supported.")
+      self.skipTest("Pruning with fp16 is currently not supported.")
+
+    if "selective_clustering" in extra_flags:
+      self.skipTest("Clustering with fp16 is currently not supported.")
 
     integration.run_synthetic(
         main=resnet_imagenet_main.run,