Fix quantization tests on opensource and run tests on TF 2.X.

alanchiao · tensorflower-gardener · commit 256658f9e14f · 2019-12-15T23:37:56.000-08:00
Tests that require conversion to TFLite don't pass in 1.XX and 2.XX yet.

PiperOrigin-RevId: 285711327
diff --git a/pip_pkg.sh b/pip_pkg.sh
@@ -40,7 +40,7 @@ cd bazel-bin/pip_pkg.runfiles/tensorflow_model_optimization
 # specifies the output dir) to setup.py, e.g.,
 #  ./pip_pkg /tmp/tensorflow_model_optimization_pkg --release
 # passes `--release` to setup.py.
-python setup.py bdist_wheel --universal ${@:2} --dist-dir="$DEST" # >/dev/null
+python3 setup.py bdist_wheel --universal ${@:2} --dist-dir="$DEST" # >/dev/null
 
 set +x
 echo -e "\nBuild complete. Wheel files are in $DEST"
diff --git a/tensorflow_model_optimization/python/core/quantization/keras/BUILD b/tensorflow_model_optimization/python/core/quantization/keras/BUILD
@@ -73,6 +73,7 @@ py_test(
     visibility = ["//visibility:public"],
     deps = [
         ":quantizers",
+        # absl/testing:parameterized dep1,
         # numpy dep1,
         # tensorflow dep1,
         # python:dtypes tensorflow dep2,
@@ -147,6 +148,7 @@ py_test(
     deps = [
         ":quantize_annotate",
         ":quantize_provider",
+        # numpy dep1,
         # tensorflow dep1,
         # python/keras tensorflow dep2,
     ],
@@ -177,6 +179,7 @@ py_test(
     deps = [
         ":quantize_aware_activation",
         ":quantizers",
+        # numpy dep1,
         # tensorflow dep1,
         # python/keras tensorflow dep2,
     ],
@@ -209,6 +212,7 @@ py_test(
     deps = [
         ":quantize_aware_activation",
         ":quantize_wrapper",
+        # absl/testing:parameterized dep1,
         # numpy dep1,
         # tensorflow dep1,
         # python/keras tensorflow dep2,
@@ -244,6 +248,7 @@ py_test(
         ":quantize",
         ":quantize_wrapper",
         # absl/testing:parameterized dep1,
+        # numpy dep1,
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/quantization/keras/tflite:tflite_quantize_registry",
     ],
@@ -256,6 +261,7 @@ py_test(
     deps = [
         ":quantize",
         ":utils",
+        # absl/testing:parameterized dep1,
         # numpy dep1,
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/keras:test_utils",
@@ -271,6 +277,7 @@ py_test(
     deps = [
         ":quantize",
         ":utils",
+        # absl/testing:parameterized dep1,
         # numpy dep1,
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/keras:test_utils",
diff --git a/tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/BUILD b/tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/BUILD
@@ -55,6 +55,7 @@ py_test(
     visibility = ["//visibility:public"],
     deps = [
         ":model_transformer",
+        # numpy dep1,
         # tensorflow dep1,
         # python/keras tensorflow dep2,
     ],
diff --git a/tensorflow_model_optimization/python/core/quantization/keras/layers/BUILD b/tensorflow_model_optimization/python/core/quantization/keras/layers/BUILD
@@ -29,6 +29,7 @@ py_test(
     deps = [
         ":conv_batchnorm",
         ":conv_batchnorm_test_utils",
+        # numpy dep1,
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/quantization/keras:quantize",
         "//tensorflow_model_optimization/python/core/quantization/keras:utils",
diff --git a/tensorflow_model_optimization/python/core/quantization/keras/layers/conv_batchnorm.py b/tensorflow_model_optimization/python/core/quantization/keras/layers/conv_batchnorm.py
@@ -16,7 +16,6 @@
 
 from __future__ import absolute_import
 from __future__ import division
-from __future__ import google_type_annotations
 from __future__ import print_function
 
 import tensorflow as tf
diff --git a/tensorflow_model_optimization/python/core/quantization/keras/layers/conv_batchnorm_test.py b/tensorflow_model_optimization/python/core/quantization/keras/layers/conv_batchnorm_test.py
@@ -75,7 +75,6 @@ def _get_asymmetric_quant_params(real_min, real_max, quant_min, quant_max):
   def _test_equal_tf_and_tflite_outputs(self,
                                         tf_model,
                                         is_tflite_quantized=False):
-    _, keras_file = tempfile.mkstemp('.h5')
     _, tflite_file = tempfile.mkstemp('.tflite')
 
     batched_input_shape = self._get_batched_input_shape()
@@ -115,10 +114,9 @@ def _test_equal_tf_and_tflite_outputs(self,
     tf_out = tf_model.predict(inp)
 
     # TensorFlow Lite inference.
-    tf.keras.models.save_model(tf_model, keras_file)
     with quantize.quantize_scope():
       utils.convert_keras_to_tflite(
-          keras_file,
+          tf_model,
           tflite_file,
           custom_objects={
               '_ConvBatchNorm2D': _ConvBatchNorm2D,
diff --git a/tensorflow_model_optimization/python/core/quantization/keras/quantize_functional_test.py b/tensorflow_model_optimization/python/core/quantization/keras/quantize_functional_test.py
@@ -22,7 +22,6 @@
 
 from absl.testing import parameterized
 
-from tensorflow.python import keras
 from tensorflow.python.platform import test
 
 from tensorflow_model_optimization.python.core.keras.testing import test_utils_mnist
@@ -54,13 +53,11 @@ def testQuantizesMnist(self):
 
     self.assertGreater(quantized_model_accuracy, 0.6)
 
-    _, quantized_keras_file = tempfile.mkstemp('.h5')
     _, quantized_tflite_file = tempfile.mkstemp('.tflite')
 
-    keras.models.save_model(quantized_model, quantized_keras_file)
     with quantize.quantize_scope():
       test_utils.convert_keras_to_tflite(
-          model_path=quantized_keras_file,
+          model=quantized_model,
           output_path=quantized_tflite_file,
           is_quantized=True)
     quantized_model_tflite_accuracy = test_utils_mnist.eval_tflite(
diff --git a/tensorflow_model_optimization/python/core/quantization/keras/quantize_integration_test.py b/tensorflow_model_optimization/python/core/quantization/keras/quantize_integration_test.py
@@ -72,7 +72,7 @@ def testSerialization(self):
         loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy'])
     quantized_model.fit(
         np.random.rand(20, 10),
-        keras.utils.np_utils.to_categorical(
+        tf.keras.utils.to_categorical(
             np.random.randint(5, size=(20, 1)), 5),
         batch_size=20)
 
@@ -91,13 +91,10 @@ def testProductionModelConversionToTFLite(self):
     annotated = quantize_annotate(model)
     quantized_model = quantize_apply(annotated)
 
-    _, keras_file = tempfile.mkstemp('.h5')
     _, tflite_file = tempfile.mkstemp('.h5')
 
-    keras.models.save_model(quantized_model, keras_file)
-
     with quantize.quantize_scope():
-      utils.convert_keras_to_tflite(keras_file, tflite_file)
+      utils.convert_keras_to_tflite(quantized_model, tflite_file)
 
 
 if __name__ == '__main__':
diff --git a/tensorflow_model_optimization/python/core/quantization/keras/quantize_wrapper_test.py b/tensorflow_model_optimization/python/core/quantization/keras/quantize_wrapper_test.py
@@ -116,7 +116,7 @@ def _get_quantized_weights(shape, dtype):  # pylint: disable=unused-argument
       assert tuple(shape) == self.weights.shape
 
       # Default values used in TFLiteRegistry.
-      return tf.fake_quant_with_min_max_vars(
+      return tf.quantization.fake_quant_with_min_max_vars(
           self.weights, -6.0, 6.0, num_bits=8, narrow_range=True)
 
     layer = layer_type(kernel_initializer=_get_random_weights, **kwargs)
@@ -141,7 +141,7 @@ def _get_quantized_weights(shape, dtype):  # pylint: disable=unused-argument
     inputs = np.random.rand(1, *input_shape)
     # `quantized_model` should apply FakeQuant. Explicitly applying to the
     # results of `model` to verify QuantizeWrapper works as expected.
-    expected_output = tf.fake_quant_with_min_max_vars(
+    expected_output = tf.quantization.fake_quant_with_min_max_vars(
         model.predict(inputs), -6.0, 6.0, num_bits=8, narrow_range=False)
     self.assertAllClose(expected_output, quantized_model.predict(inputs))
 
@@ -157,7 +157,7 @@ def testQuantizesOutputsFromLayer(self):
     model = keras.Sequential([layers.ReLU()])
 
     inputs = np.random.rand(1, 2, 1)
-    expected_output = tf.fake_quant_with_min_max_vars(
+    expected_output = tf.quantization.fake_quant_with_min_max_vars(
         model.predict(inputs), -6.0, 6.0, num_bits=8, narrow_range=False)
     self.assertAllClose(expected_output, quantized_model.predict(inputs))
 
diff --git a/tensorflow_model_optimization/python/core/quantization/keras/tflite/BUILD b/tensorflow_model_optimization/python/core/quantization/keras/tflite/BUILD
@@ -62,6 +62,7 @@ py_test(
     visibility = ["//visibility:public"],
     deps = [
         ":tflite_quantize_registry",
+        # numpy dep1,
         # tensorflow dep1,
         # python/keras tensorflow dep2,
     ],
@@ -92,6 +93,7 @@ py_test(
     visibility = ["//visibility:public"],
     deps = [
         ":tflite_transforms",
+        # numpy dep1,
         # python/keras tensorflow dep2,
         "//tensorflow_model_optimization/python/core/quantization/keras:quantize",
         "//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:model_transformer",
diff --git a/tensorflow_model_optimization/python/core/quantization/keras/tflite/tflite_quantize_registry_test.py b/tensorflow_model_optimization/python/core/quantization/keras/tflite/tflite_quantize_registry_test.py
@@ -19,6 +19,7 @@
 from __future__ import print_function
 
 import numpy as np
+import tensorflow as tf
 
 from tensorflow.python import keras
 from tensorflow.python.keras import backend as K
@@ -64,6 +65,16 @@ def _assert_activation_quantizers(self, quantizer_list):
     for quantizer in quantizer_list:
       self.assertIsInstance(quantizer, quantizers.MovingAverageQuantizer)
 
+  def _assert_kernel_equality(self, a, b):
+    # TODO(tfmot): branch on eager mode also.
+    if tf.__version__[0] == '1':
+      self.assertEqual(a, b)
+    # TODO(tfmot): find why self.assertAllEqual(a, b) doesn't work as
+    # https://www.tensorflow.org/api_docs/python/tf/test/TestCase#assertAllClose
+    # suggests).
+    else:
+      self.assertAllEqual(a.numpy(), b.numpy())
+
 
 class TFLiteQuantizeRegistryTest(test.TestCase, _TestHelper):
 
@@ -138,7 +149,7 @@ def testReturnsProvider_KerasLayer(self):
     quantize_provider.set_quantize_weights(layer, [quantize_kernel])
     quantize_provider.set_quantize_activations(layer, [quantize_activation])
 
-    self.assertEqual(layer.kernel, quantize_kernel)
+    self._assert_kernel_equality(layer.kernel, quantize_kernel)
     self.assertEqual(layer.activation, quantize_activation)
 
   def testReturnsProvider_LayerWithResultQuantizer(self):
@@ -243,7 +254,7 @@ def testSetsQuantizeWeights(self):
         ['kernel'], ['activation'], False)
     quantize_provider.set_quantize_weights(layer, [quantize_kernel])
 
-    self.assertEqual(layer.kernel, quantize_kernel)
+    self._assert_kernel_equality(layer.kernel, quantize_kernel)
 
   def testSetsQuantizeActivations(self):
     layer = self._simple_dense_layer()
diff --git a/tensorflow_model_optimization/python/core/quantization/keras/utils.py b/tensorflow_model_optimization/python/core/quantization/keras/utils.py
@@ -15,25 +15,37 @@
 # pylint: disable=protected-access
 """Quantization specific utilities for generating, saving, testing, and evaluating models."""
 
+import tempfile
+
 import tensorflow as tf
 
 
-def convert_keras_to_tflite(model_path,
+def convert_keras_to_tflite(model,
                             output_path,
                             custom_objects=None,
                             is_quantized=True):
   """Convert Keras model to TFLite."""
   if custom_objects is None:
     custom_objects = {}
 
-  converter = tf.lite.TFLiteConverter.from_keras_model_file(
-      model_path, custom_objects=custom_objects)
+  if tf.__version__[0] == '1':
+    _, keras_file = tempfile.mkstemp('.h5')
+    tf.keras.models.save_model(model, keras_file)
+    converter = tf.lite.TFLiteConverter.from_keras_model_file(
+        keras_file, custom_objects=custom_objects)
+  else:
+    converter = tf.lite.TFLiteConverter.from_keras_model(model)
   converter.experimental_new_converter = True
 
   if is_quantized:
-    converter.inference_type = tf.lite.constants.INT8
-    converter.inference_input_type = tf.lite.constants.INT8
+    if tf.__version__[0] == '1':
+      converter.inference_type = tf.lite.constants.INT8
+      converter.inference_input_type = tf.lite.constants.INT8
+    else:
+      converter.inference_type = tf.int8
+      converter.inference_input_type = tf.int8
 
+    # TODO(tfmot): API not supported in TF 2.XX.
     input_arrays = converter.get_input_arrays()
     converter.quantized_input_stats = {
         input_arrays[0]: (-128., 255.)
