Serialize/Deserialize implementation for QuantizeWrapper

PiperOrigin-RevId: 264685336

Author: nutsiepully

tensorflow_model_optimization/python/core/quantization/keras/BUILD

@@ -153,6 +153,7 @@ py_test(
     srcs_version = "PY2AND3",
     visibility = ["//visibility:public"],
     deps = [
+        ":quantize_aware_activation",
         ":quantize_wrapper",
         # numpy dep1,
         # tensorflow dep1,

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

@@ -76,7 +76,12 @@ def _should_pre_quantize(self):
     # Whether we apply quantize operations around activations depends on the
     # implementation of the specific kernel. For example, ReLUs are fused in
     # whereas Softmax ops are not. Should linear have post-quantize?
-    return self.activation.__name__ in self._PRE_ACTIVATION_TYPES
+
+    # For custom quantizations unknown in keras, we default to post
+    # quantization.
+
+    return (hasattr(self.activation, '__name__') and
+            self.activation.__name__ in self._PRE_ACTIVATION_TYPES)
 
   def _add_range_weights(self, name):
     min_var = self.quantize_wrapper.add_weight(
@@ -130,3 +135,19 @@ def quantizer_fn(x=x,
                           self._max_post_activation))
 
     return x
+
+  # `QuantizeAwareActivation` wraps the activation within a layer to perform
+  # quantization. In the process, the layer's activation is replaced with
+  # `QuantizeAwareActivation`.
+  # However, when the layer is serialized and deserialized, we want the original
+  # activation to be reconstructed. This ensures that when `QuantizeWrapper`
+  # wraps the layer, it can again replace the original activation.
+
+  @classmethod
+  def from_config(cls, config):
+    return activations.deserialize(config['activation'])
+
+  def get_config(self):
+    return {
+        'activation': activations.serialize(self.activation)
+    }

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

@@ -23,6 +23,8 @@
 from tensorflow.python import keras
 from tensorflow.python.keras import activations
 from tensorflow.python.keras import backend as K
+from tensorflow.python.keras.utils.generic_utils import deserialize_keras_object
+from tensorflow.python.keras.utils.generic_utils import serialize_keras_object
 from tensorflow.python.platform import test
 
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_aware_activation
@@ -64,8 +66,8 @@ def testAppliesQuantizationPostActivation(self):
     # 256 buckets.
     # Derived using `tf.fake_quant_with_min_max_vars`
     expected_activation = np.array(
-        [0.0, 0.0, 0.0, 0.04705906, 0.09411764, 3.011765, 5.9764705]
-    ).reshape(7, 1)
+        [0.0, 0.0, 0.0, 0.04705906, 0.09411764, 3.011765,
+         5.9764705]).reshape(7, 1)
 
     self.assertAllClose(expected_activation, model.predict(x))
 
@@ -88,6 +90,66 @@ def testAppliesQuantizationPreAndPostActivation(self):
 
     self.assertAllClose(expected_activation, model.predict(x))
 
+  def testSerializationReturnsWrappedActivation_BuiltInActivation(self):
+    activation = activations.get('tanh')
+    quantize_activation = QuantizeAwareActivation(
+        activation, self.quantizer, 0, self.TestLayer())
+
+    expected_config = {
+        'class_name': 'QuantizeAwareActivation',
+        'config': {'activation': 'tanh'}
+    }
+    serialized_quantize_activation = serialize_keras_object(quantize_activation)
+
+    self.assertEqual(expected_config, serialized_quantize_activation)
+
+    deserialized_activation = deserialize_keras_object(
+        serialized_quantize_activation,
+        custom_objects={'QuantizeAwareActivation': QuantizeAwareActivation})
+
+    self.assertEqual(activation, deserialized_activation)
+
+  def testSerializationReturnsWrappedActivation_CustomActivation(self):
+    class CustomActivation(object):
+
+      def __init__(self, key):
+        self.key = key
+
+      def get_config(self):
+        return {'key': self.key}
+
+      def __call__(self, *args, **kwargs):
+        return None
+
+      def __eq__(self, other):
+        return self.key == other.key
+
+    activation = CustomActivation('value')
+    quantize_activation = QuantizeAwareActivation(
+        activation, self.quantizer, 0, self.TestLayer())
+
+    expected_config = {
+        'class_name': 'QuantizeAwareActivation',
+        'config': {
+            'activation': {
+                'class_name': 'CustomActivation',
+                'config': {'key': 'value'}
+            }
+        }
+    }
+    serialized_quantize_activation = serialize_keras_object(quantize_activation)
+
+    self.assertEqual(expected_config, serialized_quantize_activation)
+
+    deserialized_activation = deserialize_keras_object(
+        serialized_quantize_activation,
+        custom_objects={
+            'QuantizeAwareActivation': QuantizeAwareActivation,
+            'CustomActivation': CustomActivation
+        })
+
+    self.assertEqual(activation, deserialized_activation)
+
 
 if __name__ == '__main__':
   test.main()

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

@@ -29,11 +29,13 @@
 from tensorflow.python.framework import dtypes
 from tensorflow.python.keras import backend as K
 from tensorflow.python.keras import initializers
+from tensorflow.python.keras.layers import deserialize as deserialize_layer
 from tensorflow.python.keras.layers.wrappers import Wrapper
 from tensorflow.python.keras.utils import tf_utils
+from tensorflow.python.keras.utils.generic_utils import deserialize_keras_object
+from tensorflow.python.keras.utils.generic_utils import serialize_keras_object
 
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_aware_activation
-from tensorflow_model_optimization.python.core.quantization.keras import quantize_provider as quantize_provider_mod
 
 
 class QuantizeWrapper(Wrapper):
@@ -155,29 +157,27 @@ def quantizer_fn(unquantized_weight=unquantized_weight,
 
   def get_config(self):
     base_config = super(QuantizeWrapper, self).get_config()
-    config = {'quantize_provider': self.quantize_provider}
+    config = {
+        'quantize_provider': serialize_keras_object(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,
+    # QuantizeWrapper may be constructed with any QuantizeProvider and the
+    # wrapper itself cannot know all the possible provider classes.
+    # The deserialization code should ensure the QuantizeProvider is in keras
+    # serialization scope.
+    quantize_provider = deserialize_keras_object(
+        config.pop('quantize_provider'),
         module_objects=globals(),
-        custom_objects=custom_objects)
+        custom_objects=None)
 
-    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
 
-    return cls(**config)
+    return cls(layer=layer, quantize_provider=quantize_provider, **config)
 
   @property
   def trainable(self):

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

@@ -25,11 +25,15 @@
 
 from tensorflow.python import keras
 from tensorflow.python.keras import layers
+from tensorflow.python.keras.layers import deserialize as deserialize_layer
+from tensorflow.python.keras.layers import serialize as serialize_layer
 from tensorflow.python.platform import test
 
+from tensorflow_model_optimization.python.core.quantization.keras import quantize_aware_activation
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_wrapper
 from tensorflow_model_optimization.python.core.quantization.keras.tflite import tflite_quantize_registry
 
+QuantizeAwareActivation = quantize_aware_activation.QuantizeAwareActivation
 QuantizeWrapper = quantize_wrapper.QuantizeWrapper
 TFLiteQuantizeRegistry = tflite_quantize_registry.TFLiteQuantizeRegistry
 
@@ -140,6 +144,28 @@ def _get_quantized_weights(shape, dtype):  # pylint: disable=unused-argument
         model.predict(inputs), -6.0, 6.0, num_bits=8, narrow_range=False)
     self.assertAllClose(expected_output, quantized_model.predict(inputs))
 
+  def testSerializationQuantizeWrapper(self):
+    input_shape = (2,)
+    layer = keras.layers.Dense(3)
+    wrapper = QuantizeWrapper(
+        layer=layer,
+        quantize_provider=self.quantize_registry.get_quantize_provider(layer),
+        input_shape=input_shape)
+
+    custom_objects = {
+        'QuantizeAwareActivation': QuantizeAwareActivation,
+        'QuantizeWrapper': QuantizeWrapper
+    }
+    custom_objects.update(tflite_quantize_registry._types_dict())
+
+    serialized_wrapper = serialize_layer(wrapper)
+    with keras.utils.custom_object_scope(custom_objects):
+      wrapper_from_config = deserialize_layer(serialized_wrapper)
+
+    self.assertEqual(wrapper_from_config.get_config(), wrapper.get_config())
+
+  # TODO(pulkitb): Add test to ensure weights are also preserved.
+
 
 if __name__ == '__main__':
   test.main()