Adding serialize/deserialize to Quantizers

PiperOrigin-RevId: 264460013

Author: nutsiepully

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

@@ -53,6 +53,22 @@ def __call__(self, inputs, step, training, **kwargs):
     Returns: quantized tensor.
     """
 
+  @abc.abstractmethod
+  def get_config(self):
+    raise NotImplementedError('Quantizer should implement get_config().')
+
+  @classmethod
+  def from_config(cls, config):
+    """Instantiates a `Quantizer` from its config.
+
+    Args:
+        config: Output of `get_config()`.
+
+    Returns:
+        A `Quantizer` instance.
+    """
+    return cls(**config)
+
 
 class LastValueQuantizer(Quantizer):
   """Quantize tensor based on range the last batch of values."""
@@ -96,6 +112,24 @@ def __call__(self, inputs, step, training, **kwargs):
         # TODO(pulkitb): Figure out a clean way to use name_prefix here.
     )
 
+  def get_config(self):
+    return {
+        'num_bits': self.num_bits,
+        'per_axis': self.per_axis,
+        'symmetric': self.symmetric,
+    }
+
+  def __eq__(self, other):
+    if not isinstance(other, LastValueQuantizer):
+      return False
+
+    return (self.num_bits == other.num_bits and
+            self.per_axis == other.per_axis and
+            self.symmetric == other.symmetric)
+
+  def __ne__(self, other):
+    return not self.__eq__(other)
+
 
 class MovingAverageQuantizer(Quantizer):
   """Quantize tensor based on a moving average of values across batches."""
@@ -137,3 +171,28 @@ def __call__(self, inputs, step, training, **kwargs):
         narrow_range=False,
         # TODO(pulkitb): Figure out a clean way to use name_prefix here.
     )
+
+  def get_config(self):
+    return {
+        'num_bits': self.num_bits,
+        'per_axis': self.per_axis,
+        'symmetric': self.symmetric,
+    }
+
+  def __eq__(self, other):
+    if not isinstance(other, MovingAverageQuantizer):
+      return False
+
+    return (self.num_bits == other.num_bits and
+            self.per_axis == other.per_axis and
+            self.symmetric == other.symmetric)
+
+  def __ne__(self, other):
+    return not self.__eq__(other)
+
+
+def _types_dict():
+  return {
+      'LastValueQuantizer': LastValueQuantizer,
+      'MovingAverageQuantizer': MovingAverageQuantizer
+  }

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

@@ -18,19 +18,25 @@
 from __future__ import division
 from __future__ import print_function
 
+from absl.testing import parameterized
+
 import numpy as np
 
 from tensorflow.python.client import session
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
+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.ops import variable_scope
 from tensorflow.python.ops import variables
 from tensorflow.python.platform import test
 
 from tensorflow_model_optimization.python.core.quantization.keras import quantizers
 
 
-class QuantizersTest(test.TestCase):
+@parameterized.parameters(
+    quantizers.LastValueQuantizer, quantizers.MovingAverageQuantizer)
+class QuantizersTest(test.TestCase, parameterized.TestCase):
 
   def setUp(self):
     super(QuantizersTest, self).setUp()
@@ -63,15 +69,32 @@ def _test_quantizer(quantizer):
       print('min_var: ', min_max_values[0])
       print('max_var: ', min_max_values[1])
 
-  def testLastValueQuantizer(self):
-    quantizer = quantizers.LastValueQuantizer(**self.quant_params)
+  def testQuantizer(self, quantizer_type):
+    quantizer = quantizer_type(**self.quant_params)
 
     self._test_quantizer(quantizer)
 
-  def testMovingAverageQuantizer(self):
-    quantizer = quantizers.MovingAverageQuantizer(**self.quant_params)
+  def testSerialization(self, quantizer_type):
+    quantizer = quantizer_type(**self.quant_params)
 
-    self._test_quantizer(quantizer)
+    expected_config = {
+        'class_name': quantizer_type.__name__,
+        'config': {
+            'num_bits': 8,
+            'per_axis': False,
+            'symmetric': False
+        }
+    }
+    serialized_quantizer = serialize_keras_object(quantizer)
+
+    self.assertEqual(expected_config, serialized_quantizer)
+
+    quantizer_from_config = deserialize_keras_object(
+        serialized_quantizer,
+        module_objects=globals(),
+        custom_objects=quantizers._types_dict())
+
+    self.assertEqual(quantizer, quantizer_from_config)
 
 
 if __name__ == '__main__':