QuantizeAnnotate add registry support.

PiperOrigin-RevId: 255683228

Author: nutsiepully

tensorflow_model_optimization/python/core/quantization/keras/BUILD

@@ -107,6 +107,8 @@ py_library(
     srcs_version = "PY2AND3",
     visibility = ["//visibility:public"],
     deps = [
+        ":quantize_emulatable_layer",
+        ":quantize_emulate_registry",
         # tensorflow dep1,
         # python/keras tensorflow dep2,
     ],
@@ -122,6 +124,7 @@ py_test(
     visibility = ["//visibility:public"],
     deps = [
         ":quantize_annotate",
+        ":quantize_emulatable_layer",
         # tensorflow dep1,
         # python/keras tensorflow dep2,
     ],

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

@@ -20,6 +20,11 @@
 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
 
 
 class QuantizeAnnotate(Wrapper):
@@ -35,6 +40,10 @@ class QuantizeAnnotate(Wrapper):
   modified.
   """
 
+  _UNSUPPORTED_LAYER_ERROR_MSG = (
+      'Layer {} not supported for quantization. Layer should either inherit '
+      'QuantizeEmulatableLayer or be a supported keras built-in layer.')
+
   def __init__(self,
                layer,
                num_bits,
@@ -52,6 +61,12 @@ def __init__(self,
         the minimum and maximum of each quantization range separately.
       **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

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

@@ -24,20 +24,58 @@
 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
 
 
 class QuantizeAnnotateTest(test.TestCase):
 
-  def testAppliesWrapperToAllClasses(self):
-    layer = keras.layers.Dense(5, activation='relu', input_shape=(10,))
+  def setUp(self):
+    self.quant_params = {
+        'num_bits': 8,
+        'narrow_range': True,
+        'symmetric': True
+    }
+
+  def testRaisesErrorForUnsupportedLayer(self):
+    class CustomLayer(keras.layers.Dense):
+      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 testAnnotatesKerasLayer(self):
+    layer = keras.layers.Dense(5, activation='relu', input_shape=(10,))
     model = keras.Sequential([layer])
-    wrapped_model = keras.Sequential([
+
+    annotated_model = keras.Sequential([
         quantize_annotate.QuantizeAnnotate(
-            layer, num_bits=8, input_shape=(10,))])
+            layer, input_shape=(10,), **self.quant_params)])
+
+    annotated_layer = annotated_model.layers[0]
+    self.assertIsInstance(annotated_layer.layer, keras.layers.Dense)
+    self.assertEqual(
+        self.quant_params, annotated_layer.get_quantize_params())
 
+    # Annotated model should not affect computation. Returns same results.
     x_test = np.random.rand(10, 10)
-    self.assertAllEqual(model.predict(x_test), wrapped_model.predict(x_test))
+    self.assertAllEqual(model.predict(x_test), annotated_model.predict(x_test))
 
 if __name__ == '__main__':
   test.main()