Serialize/Deserialize support for QuantizeProviders

PiperOrigin-RevId: 264463346

Author: nutsiepully

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

@@ -43,6 +43,9 @@ def set_quantize_weights(self, layer, quantize_weights):
     def set_quantize_activations(self, layer, quantize_activations):
       pass
 
+    def get_config(self):
+      pass
+
   def testAnnotatesKerasLayer(self):
     layer = keras.layers.Dense(5, activation='relu', input_shape=(10,))
     model = keras.Sequential([layer])

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

@@ -117,6 +117,9 @@ def set_quantize_weights(self, layer, quantize_weights):
       def set_quantize_activations(self, layer, quantize_activations):
         pass
 
+      def get_config(self):
+        pass
+
     quantize_provider = TestQuantizeProvider()
 
     model = keras.Sequential([

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

@@ -76,3 +76,7 @@ def set_quantize_activations(self, layer, quantize_activations):
       quantizer.
     """
     raise NotImplementedError('Must be implemented in subclasses.')
+
+  @abc.abstractmethod
+  def get_config(self):
+    raise NotImplementedError('QuantizeProvider should implement get_config().')

tensorflow_model_optimization/python/core/quantization/keras/tflite/tflite_quantize_registry.py

@@ -333,6 +333,39 @@ def set_quantize_activations(self, layer, quantize_activations):
         zip(self.activation_attrs, quantize_activations):
       setattr(layer, activation_attr, activation)
 
+  @classmethod
+  def from_config(cls, config):
+    """Instantiates a `TFLiteQuantizeProvider` from its config.
+
+    Args:
+        config: Output of `get_config()`.
+
+    Returns:
+        A `TFLiteQuantizeProvider` instance.
+    """
+    return cls(**config)
+
+  def get_config(self):
+    # TODO(pulkitb): Add weight and activation quantizer to config.
+    # Currently it's created internally, but ideally the quantizers should be
+    # part of the constructor and passed in from the registry.
+    return {
+        'weight_attrs': self.weight_attrs,
+        'activation_attrs': self.activation_attrs,
+    }
+
+  def __eq__(self, other):
+    if not isinstance(other, TFLiteQuantizeProvider):
+      return False
+
+    return (self.weight_attrs == other.weight_attrs and
+            self.activation_attrs == self.activation_attrs and
+            self.weight_quantizer == other.weight_quantizer and
+            self.activation_quantizer == other.activation_quantizer)
+
+  def __ne__(self, other):
+    return not self.__eq__(other)
+
 
 class TFLiteQuantizeProviderRNN(TFLiteQuantizeProvider, _RNNHelper):
   """QuantizeProvider for RNN layers."""
@@ -402,3 +435,10 @@ def set_quantize_activations(self, layer, quantize_activations):
       for activation_attr in activation_attrs_cell:
         setattr(rnn_cell, activation_attr, quantize_activations[i])
         i += 1
+
+
+def _types_dict():
+  return {
+      'TFLiteQuantizeProvider': TFLiteQuantizeProvider,
+      'TFLiteQuantizeProviderRNN': TFLiteQuantizeProviderRNN
+  }

tensorflow_model_optimization/python/core/quantization/keras/tflite/tflite_quantize_registry_test.py

@@ -23,6 +23,8 @@
 from tensorflow.python import keras
 from tensorflow.python.keras import backend as K
 from tensorflow.python.keras import layers as l
+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 quantizers
@@ -271,6 +273,28 @@ def testSetsQuantizeActivations_ErrorOnWrongNumberOfActivations(self):
       quantize_provider.set_quantize_activations(
           layer, [quantize_activation, quantize_activation])
 
+  def testSerialization(self):
+    quantize_provider = tflite_quantize_registry.TFLiteQuantizeProvider(
+        ['kernel'], ['activation'])
+
+    expected_config = {
+        'class_name': 'TFLiteQuantizeProvider',
+        'config': {
+            'weight_attrs': ['kernel'],
+            'activation_attrs': ['activation'],
+        }
+    }
+    serialized_quantize_provider = serialize_keras_object(quantize_provider)
+
+    self.assertEqual(expected_config, serialized_quantize_provider)
+
+    quantize_provider_from_config = deserialize_keras_object(
+        serialized_quantize_provider,
+        module_objects=globals(),
+        custom_objects=tflite_quantize_registry._types_dict())
+
+    self.assertEqual(quantize_provider, quantize_provider_from_config)
+
 
 class TFLiteQuantizeProviderRNNTest(test.TestCase, _TestHelper):
 
@@ -367,6 +391,28 @@ def testSetsQuantizeActivations_ErrorOnWrongNumberOfActivations(self):
       self.quantize_provider.set_quantize_activations(
           self.layer, [quantize_activation, quantize_activation])
 
+  def testSerialization(self):
+    expected_config = {
+        'class_name': 'TFLiteQuantizeProviderRNN',
+        'config': {
+            'weight_attrs': [['kernel', 'recurrent_kernel'],
+                             ['kernel', 'recurrent_kernel']],
+            'activation_attrs': [['activation', 'recurrent_activation'],
+                                 ['activation', 'recurrent_activation']],
+        }
+    }
+    serialized_quantize_provider = serialize_keras_object(
+        self.quantize_provider)
+
+    self.assertEqual(expected_config, serialized_quantize_provider)
+
+    quantize_provider_from_config = deserialize_keras_object(
+        serialized_quantize_provider,
+        module_objects=globals(),
+        custom_objects=tflite_quantize_registry._types_dict())
+
+    self.assertEqual(self.quantize_provider, quantize_provider_from_config)
+
 
 if __name__ == '__main__':
   test.main()