Reuse test_utils in compression tests.

PiperOrigin-RevId: 370829522

Author: fredrec

tensorflow_model_optimization/python/core/common/keras/compression/algorithms/BUILD

@@ -22,8 +22,8 @@ py_strict_test(
     python_version = "PY3",
     deps = [
         ":same_training_and_inference",
-        # numpy dep1,
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/keras/testing:test_utils_mnist",
     ],
 )
 
@@ -44,8 +44,8 @@ py_strict_test(
     python_version = "PY3",
     deps = [
         ":different_training_and_inference",
-        # numpy dep1,
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/keras/testing:test_utils_mnist",
     ],
 )
 
@@ -66,8 +66,8 @@ py_strict_test(
     python_version = "PY3",
     deps = [
         ":bias_only",
-        # numpy dep1,
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/keras/testing:test_utils_mnist",
     ],
 )
 
@@ -112,7 +112,7 @@ py_strict_test(
     python_version = "PY3",
     deps = [
         ":periodical_update_and_scheduling",
-        # numpy dep1,
         # tensorflow dep1,
+        "//tensorflow_model_optimization/python/core/keras/testing:test_utils_mnist",
     ],
 )

tensorflow_model_optimization/python/core/common/keras/compression/algorithms/bias_only_test.py

@@ -17,10 +17,10 @@
 import os
 import tempfile
 
-import numpy as np
 import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.common.keras.compression.algorithms import bias_only
+from tensorflow_model_optimization.python.core.keras.testing import test_utils_mnist
 
 
 def _build_model():
@@ -81,38 +81,6 @@ def _convert_to_tflite(saved_model_dir):
   return tflite_file
 
 
-# TODO(tfmot): reuse test_utils_mnist.py.
-def _test_tflite(tflite_file):
-  interpreter = tf.lite.Interpreter(model_path=tflite_file)
-  interpreter.allocate_tensors()
-
-  input_index = interpreter.get_input_details()[0]['index']
-  output_index = interpreter.get_output_details()[0]['index']
-
-  (_, _), (x_test, y_test) = _get_dataset()
-
-  # Testing the entire dataset is too slow. Verifying only 300 of 10k samples.
-  x_test = x_test[0:300, :]
-  y_test = y_test[0:300]
-
-  total_seen = 0
-  num_correct = 0
-
-  for img, label in zip(x_test, y_test):
-    batch_input_shape = (1, 28, 28)
-    inp = img.reshape(batch_input_shape)
-    inp = inp.astype(np.float32)
-    total_seen += 1
-    interpreter.set_tensor(input_index, inp)
-    interpreter.invoke()
-    predictions = interpreter.get_tensor(output_index)
-
-    if np.argmax(predictions) == label:
-      num_correct += 1
-
-  return float(num_correct) / float(total_seen)
-
-
 def _get_directory_size_in_bytes(directory):
   total = 0
   try:
@@ -169,7 +137,7 @@ def testBiasOnly_HasReasonableAccuracy_TFLite(self):
     saved_model_dir = _save_as_saved_model(compressed_model)
     compressed_tflite_file = _convert_to_tflite(saved_model_dir)
 
-    accuracy = _test_tflite(compressed_tflite_file)
+    accuracy = test_utils_mnist.eval_tflite(compressed_tflite_file)
 
     self.assertGreater(accuracy, 0.60)
 

tensorflow_model_optimization/python/core/common/keras/compression/algorithms/different_training_and_inference_test.py

@@ -16,10 +16,10 @@
 import os
 import tempfile
 
-import numpy as np
 import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.common.keras.compression.algorithms import different_training_and_inference as svd
+from tensorflow_model_optimization.python.core.keras.testing import test_utils_mnist
 
 
 # TODO(tfmot): dedup.
@@ -81,38 +81,6 @@ def _convert_to_tflite(saved_model_dir):
   return tflite_file
 
 
-# TODO(tfmot): reuse test_utils_mnist.py.
-def _test_tflite(tflite_file):
-  interpreter = tf.lite.Interpreter(model_path=tflite_file)
-  interpreter.allocate_tensors()
-
-  input_index = interpreter.get_input_details()[0]['index']
-  output_index = interpreter.get_output_details()[0]['index']
-
-  (_, _), (x_test, y_test) = _get_dataset()
-
-  # Testing the entire dataset is too slow. Verifying only 300 of 10k samples.
-  x_test = x_test[0:300, :]
-  y_test = y_test[0:300]
-
-  total_seen = 0
-  num_correct = 0
-
-  for img, label in zip(x_test, y_test):
-    batch_input_shape = (1, 28, 28)
-    inp = img.reshape(batch_input_shape)
-    inp = inp.astype(np.float32)
-    total_seen += 1
-    interpreter.set_tensor(input_index, inp)
-    interpreter.invoke()
-    predictions = interpreter.get_tensor(output_index)
-
-    if np.argmax(predictions) == label:
-      num_correct += 1
-
-  return float(num_correct) / float(total_seen)
-
-
 def _get_directory_size_in_bytes(directory):
   total = 0
   try:
@@ -192,7 +160,7 @@ def testSVD_HasReasonableAccuracy_TFLite(self):
     saved_model_dir = _save_as_saved_model(model_for_inference)
     compressed_tflite_file = _convert_to_tflite(saved_model_dir)
 
-    accuracy = _test_tflite(compressed_tflite_file)
+    accuracy = test_utils_mnist.eval_tflite(compressed_tflite_file)
 
     self.assertGreater(accuracy, 0.60)
 

tensorflow_model_optimization/python/core/common/keras/compression/algorithms/periodical_update_and_scheduling_test.py

@@ -17,10 +17,10 @@
 import os
 import tempfile
 
-import numpy as np
 import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.common.keras.compression.algorithms import periodical_update_and_scheduling as svd
+from tensorflow_model_optimization.python.core.keras.testing import test_utils_mnist
 
 
 def _build_model():
@@ -81,38 +81,6 @@ def _convert_to_tflite(saved_model_dir):
   return tflite_file
 
 
-# TODO(tfmot): reuse test_utils_mnist.py.
-def _test_tflite(tflite_file):
-  interpreter = tf.lite.Interpreter(model_path=tflite_file)
-  interpreter.allocate_tensors()
-
-  input_index = interpreter.get_input_details()[0]['index']
-  output_index = interpreter.get_output_details()[0]['index']
-
-  (_, _), (x_test, y_test) = _get_dataset()
-
-  # Testing the entire dataset is too slow. Verifying only 300 of 10k samples.
-  x_test = x_test[0:300, :]
-  y_test = y_test[0:300]
-
-  total_seen = 0
-  num_correct = 0
-
-  for img, label in zip(x_test, y_test):
-    batch_input_shape = (1, 28, 28)
-    inp = img.reshape(batch_input_shape)
-    inp = inp.astype(np.float32)
-    total_seen += 1
-    interpreter.set_tensor(input_index, inp)
-    interpreter.invoke()
-    predictions = interpreter.get_tensor(output_index)
-
-    if np.argmax(predictions) == label:
-      num_correct += 1
-
-  return float(num_correct) / float(total_seen)
-
-
 def _get_directory_size_in_bytes(directory):
   total = 0
   try:
@@ -204,7 +172,7 @@ def testSVD_HasReasonableAccuracy_TFLite(self):
     saved_model_dir = _save_as_saved_model(compressed_model)
     compressed_tflite_file = _convert_to_tflite(saved_model_dir)
 
-    accuracy = _test_tflite(compressed_tflite_file)
+    accuracy = test_utils_mnist.eval_tflite(compressed_tflite_file)
 
     self.assertGreater(accuracy, 0.60)
 

tensorflow_model_optimization/python/core/common/keras/compression/algorithms/same_training_and_inference_test.py

@@ -17,10 +17,10 @@
 import os
 import tempfile
 
-import numpy as np
 import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.common.keras.compression.algorithms import same_training_and_inference as svd
+from tensorflow_model_optimization.python.core.keras.testing import test_utils_mnist
 
 
 def _build_model():
@@ -81,38 +81,6 @@ def _convert_to_tflite(saved_model_dir):
   return tflite_file
 
 
-# TODO(tfmot): reuse test_utils_mnist.py.
-def _test_tflite(tflite_file):
-  interpreter = tf.lite.Interpreter(model_path=tflite_file)
-  interpreter.allocate_tensors()
-
-  input_index = interpreter.get_input_details()[0]['index']
-  output_index = interpreter.get_output_details()[0]['index']
-
-  (_, _), (x_test, y_test) = _get_dataset()
-
-  # Testing the entire dataset is too slow. Verifying only 300 of 10k samples.
-  x_test = x_test[0:300, :]
-  y_test = y_test[0:300]
-
-  total_seen = 0
-  num_correct = 0
-
-  for img, label in zip(x_test, y_test):
-    batch_input_shape = (1, 28, 28)
-    inp = img.reshape(batch_input_shape)
-    inp = inp.astype(np.float32)
-    total_seen += 1
-    interpreter.set_tensor(input_index, inp)
-    interpreter.invoke()
-    predictions = interpreter.get_tensor(output_index)
-
-    if np.argmax(predictions) == label:
-      num_correct += 1
-
-  return float(num_correct) / float(total_seen)
-
-
 def _get_directory_size_in_bytes(directory):
   total = 0
   try:
@@ -194,7 +162,7 @@ def testSVD_HasReasonableAccuracy_TFLite(self):
     saved_model_dir = _save_as_saved_model(compressed_model)
     compressed_tflite_file = _convert_to_tflite(saved_model_dir)
 
-    accuracy = _test_tflite(compressed_tflite_file)
+    accuracy = test_utils_mnist.eval_tflite(compressed_tflite_file)
 
     self.assertGreater(accuracy, 0.60)
 

tensorflow_model_optimization/python/core/keras/testing/test_utils_mnist.py

@@ -24,7 +24,7 @@
 def layers_list():
   return [
       l.Conv2D(32, 5, padding='same', activation='relu',
-               input_shape=input_shape()),
+               input_shape=image_input_shape()),
       l.MaxPooling2D((2, 2), (2, 2), padding='same'),
       # TODO(pulkitb): Add BatchNorm when transformations are ready.
       # l.BatchNormalization(),
@@ -43,7 +43,7 @@ def sequential_model():
 
 def functional_model():
   """Builds an MNIST functional model."""
-  inp = keras.Input(input_shape())
+  inp = keras.Input(image_input_shape())
   x = l.Conv2D(32, 5, padding='same', activation='relu')(inp)
   x = l.MaxPooling2D((2, 2), (2, 2), padding='same')(x)
   # TODO(pulkitb): Add BatchNorm when transformations are ready.
@@ -58,7 +58,7 @@ def functional_model():
   return keras.models.Model([inp], [out])
 
 
-def input_shape(img_rows=28, img_cols=28):
+def image_input_shape(img_rows=28, img_cols=28):
   if tf.keras.backend.image_data_format() == 'channels_first':
     return 1, img_rows, img_cols
   else:
@@ -92,9 +92,11 @@ def preprocessed_data(img_rows=28,
 
 def eval_tflite(model_path):
   """Evaluate mnist in TFLite for accuracy."""
+
   interpreter = tf.lite.Interpreter(model_path=model_path)
   interpreter.allocate_tensors()
   input_index = interpreter.get_input_details()[0]['index']
+  input_shape = interpreter.get_input_details()[0]['shape']
   output_index = interpreter.get_output_details()[0]['index']
 
   _, _, x_test, y_test = preprocessed_data()
@@ -106,8 +108,7 @@ def eval_tflite(model_path):
   num_correct = 0
 
   for img, label in zip(x_test, y_test):
-    batch_input_shape = (1,) + input_shape()
-    inp = img.reshape(batch_input_shape)
+    inp = img.reshape(input_shape)
     total_seen += 1
     interpreter.set_tensor(input_index, inp)
     interpreter.invoke()