Move quantize_annotate to new API structure.

This CL begins the process to migrate to the new API
structure. Changing in entirety would break lots of
tests, or require a large CL changing the entire code.

This CL only changes the quantize_annotate layers.

PiperOrigin-RevId: 263236324

Author: nutsiepully

tensorflow_model_optimization/python/core/quantization/keras/BUILD

@@ -183,8 +183,7 @@ py_library(
     srcs_version = "PY2AND3",
     visibility = ["//visibility:public"],
     deps = [
-        ":quantize_emulatable_layer",
-        ":quantize_emulate_registry",
+        ":quantize_provider",
         # tensorflow dep1,
         # python/keras tensorflow dep2,
     ],
@@ -200,7 +199,7 @@ py_test(
     visibility = ["//visibility:public"],
     deps = [
         ":quantize_annotate",
-        ":quantize_emulatable_layer",
+        ":quantize_provider",
         # tensorflow dep1,
         # python/keras tensorflow dep2,
     ],

tensorflow_model_optimization/python/core/quantization/keras/quantize_annotate.py

@@ -20,11 +20,7 @@
 from __future__ import print_function
 
 from tensorflow.python.keras.layers.wrappers import Wrapper
-from tensorflow_model_optimization.python.core.quantization.keras import quantize_emulatable_layer
-from tensorflow_model_optimization.python.core.quantization.keras import quantize_emulate_registry
-
-QuantizeEmulatableLayer = quantize_emulatable_layer.QuantizeEmulatableLayer
-QuantizeEmulateRegistry = quantize_emulate_registry.QuantizeEmulateRegistry
+from tensorflow_model_optimization.python.core.quantization.keras import quantize_provider as quantize_provider_mod
 
 
 class QuantizeAnnotate(Wrapper):
@@ -46,52 +42,52 @@ class QuantizeAnnotate(Wrapper):
 
   def __init__(self,
                layer,
-               num_bits,
-               narrow_range=True,
-               symmetric=True,
+               quantize_provider=None,
                **kwargs):
     """Create a quantize annotate wrapper over a keras layer.
 
     Args:
       layer: The keras layer to be quantized.
-      num_bits: Number of bits for quantization
-      narrow_range: Whether to use the narrow quantization range [1; 2^num_bits
-        - 1] or wide range [0; 2^num_bits - 1].
-      symmetric: If true, use symmetric quantization limits instead of training
-        the minimum and maximum of each quantization range separately.
+      quantize_provider: `QuantizeProvider` to quantize layer.
       **kwargs: Additional keyword arguments to be passed to the keras layer.
     """
-
-    if not isinstance(layer, QuantizeEmulatableLayer) and \
-        not QuantizeEmulateRegistry.supports(layer):
-      raise ValueError(
-          self._UNSUPPORTED_LAYER_ERROR_MSG.format(layer.__class__))
-
     super(QuantizeAnnotate, self).__init__(layer, **kwargs)
 
-    self._num_bits = num_bits
-    self._narrow_range = narrow_range
-    self._symmetric = symmetric
+    self.quantize_provider = quantize_provider
 
   def call(self, inputs, training=None):
     return self.layer.call(inputs)
 
   def get_quantize_params(self):
+    # TODO(pulkitb): Keep around function so rest of code works. Remove later.
     return {
-        'num_bits': self._num_bits,
-        'symmetric': self._symmetric,
-        'narrow_range': self._narrow_range
+        'num_bits': 8,
+        'symmetric': True,
+        'narrow_range': True
     }
 
   def get_config(self):
     base_config = super(QuantizeAnnotate, self).get_config()
-    config = self.get_quantize_params()
+    config = {
+        'quantize_provider': self.quantize_provider
+    }
     return dict(list(base_config.items()) + list(config.items()))
 
   @classmethod
   def from_config(cls, config):
     config = config.copy()
 
+    quantize_provider = config.pop('quantize_provider')
+    from tensorflow.python.keras.utils.generic_utils import deserialize_keras_object  # pylint: disable=g-import-not-at-top
+    # TODO(pulkitb): Add all known `QuantizeProvider`s to custom_objects
+    custom_objects = {
+        'QuantizeProvider': quantize_provider_mod.QuantizeProvider
+    }
+    config['quantize_provider'] = deserialize_keras_object(
+        quantize_provider,
+        module_objects=globals(),
+        custom_objects=custom_objects)
+
     from tensorflow.python.keras.layers import deserialize as deserialize_layer  # pylint: disable=g-import-not-at-top
     layer = deserialize_layer(config.pop('layer'))
     config['layer'] = layer

tensorflow_model_optimization/python/core/quantization/keras/quantize_annotate_test.py

@@ -24,58 +24,36 @@
 from tensorflow.python.platform import test
 
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_annotate
-from tensorflow_model_optimization.python.core.quantization.keras import quantize_emulatable_layer
-
-QuantizeEmulatableLayer = quantize_emulatable_layer.QuantizeEmulatableLayer
+from tensorflow_model_optimization.python.core.quantization.keras import quantize_provider as quantize_provider_mod
 
 
 class QuantizeAnnotateTest(test.TestCase):
 
-  def setUp(self):
-    self.quant_params = {
-        'num_bits': 8,
-        'narrow_range': True,
-        'symmetric': True
-    }
+  class TestQuantizeProvider(quantize_provider_mod.QuantizeProvider):
 
-  def testRaisesErrorForUnsupportedLayer(self):
-    class CustomLayer(keras.layers.Dense):
+    def get_weights_and_quantizers(self, layer):
       pass
 
-    with self.assertRaises(ValueError):
-      quantize_annotate.QuantizeAnnotate(CustomLayer(10), **self.quant_params)
-
-  def testAnnotatesCustomQuantizableLayer(self):
-    class CustomLayerQuantizable(keras.layers.Dense, QuantizeEmulatableLayer):
-      def get_quantizable_weights(self):  # pylint: disable=g-wrong-blank-lines
-        return [self.kernel]
-
-      def set_quantizable_weights(self, weights):
-        self.kernel = weights[0]
-
-    annotated_layer = quantize_annotate.QuantizeAnnotate(
-        CustomLayerQuantizable(10), **self.quant_params)
-
-    self.assertIsInstance(annotated_layer.layer, CustomLayerQuantizable)
-    self.assertEqual(
-        self.quant_params, annotated_layer.get_quantize_params())
+    def get_activations_and_quantizers(self, layer):
+      pass
 
   def testAnnotatesKerasLayer(self):
     layer = keras.layers.Dense(5, activation='relu', input_shape=(10,))
     model = keras.Sequential([layer])
 
+    quantize_provider = self.TestQuantizeProvider()
     annotated_model = keras.Sequential([
         quantize_annotate.QuantizeAnnotate(
-            layer, input_shape=(10,), **self.quant_params)])
+            layer, quantize_provider=quantize_provider, input_shape=(10,))])
 
     annotated_layer = annotated_model.layers[0]
-    self.assertIsInstance(annotated_layer.layer, keras.layers.Dense)
-    self.assertEqual(
-        self.quant_params, annotated_layer.get_quantize_params())
+    self.assertEqual(layer, annotated_layer.layer)
+    self.assertEqual(quantize_provider, annotated_layer.quantize_provider)
 
     # Annotated model should not affect computation. Returns same results.
     x_test = np.random.rand(10, 10)
     self.assertAllEqual(model.predict(x_test), annotated_model.predict(x_test))
 
+
 if __name__ == '__main__':
   test.main()

tensorflow_model_optimization/python/core/quantization/keras/quantize_emulate.py

@@ -88,12 +88,7 @@ def _QuantizeList(layers, **params):
 
 
 # TODO(pulkitb): Enable lint naming is fixed and made consistent.
-def quantize_annotate(
-    to_quantize,
-    num_bits,
-    narrow_range=True,
-    symmetric=True,
-    **kwargs):  # pylint: disable=invalid-name
+def quantize_annotate(to_quantize, **kwargs):  # pylint: disable=invalid-name
   """Specify a layer or model to be quantized.
 
   This function does not apply an quantization emulation operations. It merely
@@ -102,11 +97,6 @@ def quantize_annotate(
 
   Args:
     to_quantize: Keras layer or model to be quantized.
-    num_bits: Number of bits for quantization
-    narrow_range: Whether to use the narrow quantization range [1; 2^num_bits
-      - 1] or wide range [0; 2^num_bits - 1].
-    symmetric: If true, use symmetric quantization limits instead of training
-      the minimum and maximum of each quantization range separately.
     **kwargs: Additional keyword arguments to be passed to the keras layer.
 
   Returns:
@@ -119,20 +109,15 @@ def _add_quant_wrapper(layer):
     if isinstance(layer, quant_annotate.QuantizeAnnotate):
       return layer
 
-    return quant_annotate.QuantizeAnnotate(layer, **quant_params)
-
-  quant_params = {
-      'num_bits': num_bits,
-      'narrow_range': narrow_range,
-      'symmetric': symmetric
-  }
+    return quant_annotate.QuantizeAnnotate(layer)
 
   if isinstance(to_quantize, keras.Model):
     return keras.models.clone_model(
         to_quantize, input_tensors=None, clone_function=_add_quant_wrapper)
   elif isinstance(to_quantize, keras.layers.Layer):
-    quant_params.update(**kwargs)
-    return quant_annotate.QuantizeAnnotate(to_quantize, **quant_params)
+    # TODO(pulkitb): Since annotation for model and layer have different
+    # parameters, we should likely remove support for layers here.
+    return quant_annotate.QuantizeAnnotate(to_quantize, **kwargs)
 
 
 def quantize_apply(model):

tensorflow_model_optimization/python/core/quantization/keras/quantize_emulate_test.py

@@ -27,6 +27,7 @@
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_aware_activation
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_emulate
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_emulate_wrapper
+from tensorflow_model_optimization.python.core.quantization.keras import quantize_provider as quantize_provider_mod
 
 quantize_annotate = quantize_emulate.quantize_annotate
 QuantizeEmulate = quantize_emulate.QuantizeEmulate
@@ -67,35 +68,20 @@ def testQuantizeEmulateList(self):
 
 class QuantizeAnnotateTest(test.TestCase):
 
-  def setUp(self):
-    self.quant_params = {
-        'num_bits': 8,
-        'narrow_range': True,
-        'symmetric': True
-    }
-
-  def _assertQuantParams(self, layer, quant_params):
-    layer_params = {
-        'num_bits': layer._num_bits,
-        'narrow_range': layer._narrow_range,
-        'symmetric': layer._symmetric
-    }
-    self.assertEqual(quant_params, layer_params)
-
-  def _assertWrappedLayer(self, layer, quant_params):
+  def _assertWrappedLayer(self, layer, quantize_provider=None):
     self.assertIsInstance(layer, quant_annotate.QuantizeAnnotate)
-    self._assertQuantParams(layer, quant_params)
+    self.assertEqual(quantize_provider, layer.quantize_provider)
 
-  def _assertWrappedSequential(self, model, quant_params):
+  def _assertWrappedModel(self, model):
     for layer in model.layers:
-      self._assertWrappedLayer(layer, quant_params)
+      self._assertWrappedLayer(layer)
 
   def testQuantizeAnnotateLayer(self):
     layer = keras.layers.Dense(10, input_shape=(5,))
     wrapped_layer = quantize_annotate(
-        layer, input_shape=(5,), **self.quant_params)
+        layer, input_shape=(5,))
 
-    self._assertWrappedLayer(wrapped_layer, self.quant_params)
+    self._assertWrappedLayer(wrapped_layer)
 
     inputs = np.random.rand(1, 5)
     model = keras.Sequential([layer])
@@ -110,24 +96,33 @@ def testQuantizeAnnotateModel(self):
         keras.layers.Dense(10, input_shape=(5,)),
         keras.layers.Dropout(0.4)
     ])
-    annotated_model = quantize_annotate(model, **self.quant_params)
+    annotated_model = quantize_annotate(model)
 
-    self._assertWrappedSequential(annotated_model, self.quant_params)
+    self._assertWrappedModel(annotated_model)
 
     inputs = np.random.rand(1, 5)
     self.assertAllEqual(model.predict(inputs), annotated_model.predict(inputs))
 
   def testQuantizeAnnotateModel_HasAnnotatedLayers(self):
-    layer_params = {'num_bits': 4, 'narrow_range': False, 'symmetric': False}
+    class TestQuantizeProvider(quantize_provider_mod.QuantizeProvider):
+
+      def get_weights_and_quantizers(self, layer):
+        pass
+
+      def get_activations_and_quantizers(self, layer):
+        pass
+
+    quantize_provider = TestQuantizeProvider()
+
     model = keras.Sequential([
         keras.layers.Dense(10, input_shape=(5,)),
-        quantize_annotate(keras.layers.Dense(5), **layer_params)
+        quant_annotate.QuantizeAnnotate(
+            keras.layers.Dense(5), quantize_provider=quantize_provider)
     ])
+    annotated_model = quantize_annotate(model)
 
-    annotated_model = quantize_annotate(model, **self.quant_params)
-
-    self._assertWrappedLayer(annotated_model.layers[0], self.quant_params)
-    self._assertWrappedLayer(annotated_model.layers[1], layer_params)
+    self._assertWrappedLayer(annotated_model.layers[0])
+    self._assertWrappedLayer(annotated_model.layers[1], quantize_provider)
     # Ensure an already annotated layer is not wrapped again.
     self.assertIsInstance(annotated_model.layers[1].layer, keras.layers.Dense)
 
@@ -181,7 +176,7 @@ def testRaisesErrorNoAnnotatedLayers_Functional(self):
 
   def testRaisesErrorModelNotBuilt(self):
     model = keras.Sequential([
-        quantize_annotate(keras.layers.Dense(10), **self.quant_params1)])
+        quantize_annotate(keras.layers.Dense(10))])
 
     self.assertFalse(model.built)
     with self.assertRaises(ValueError):
@@ -191,16 +186,15 @@ def testRaisesErrorModelNotBuilt(self):
 
   def _get_annotated_sequential_model(self):
     return keras.Sequential([
-        quantize_annotate(keras.layers.Conv2D(32, 5), input_shape=(28, 28, 1),
-                          **self.quant_params1),
-        quantize_annotate(keras.layers.Dense(10), **self.quant_params2)
+        quantize_annotate(keras.layers.Conv2D(32, 5), input_shape=(28, 28, 1)),
+        quantize_annotate(keras.layers.Dense(10))
     ])
 
   def _get_annotated_functional_model(self):
     inputs = keras.Input(shape=(28, 28, 1))
     x = quantize_annotate(
-        keras.layers.Conv2D(32, 5), **self.quant_params1)(inputs)
-    results = quantize_annotate(keras.layers.Dense(10), **self.quant_params2)(x)
+        keras.layers.Conv2D(32, 5))(inputs)
+    results = quantize_annotate(keras.layers.Dense(10))(x)
 
     return keras.Model(inputs=inputs, outputs=results)
 
@@ -283,8 +277,7 @@ def testQuantizesActivationsWithinLayer_Sequential(self):
     model = keras.Sequential([
         quantize_annotate(
             keras.layers.Conv2D(32, 5, activation='relu'),
-            input_shape=(28, 28, 1),
-            **quant_params)
+            input_shape=(28, 28, 1))
     ])
 
     quantized_model = quantize_emulate.quantize_apply(model)
@@ -304,8 +297,7 @@ def testQuantizesActivationsWithinLayer_Functional(self):
 
     inputs = keras.Input(shape=(28, 28, 1))
     results = quantize_annotate(
-        keras.layers.Conv2D(32, 5, activation='relu'),
-        **self.quant_params1)(inputs)
+        keras.layers.Conv2D(32, 5, activation='relu'))(inputs)
     model = keras.Model(inputs=inputs, outputs=results)
 
     quantized_model = quantize_emulate.quantize_apply(model)