Fix an issue where multiple instances of DefaultNBitQuantizeRegistry could interfere with each other's configuration.

PiperOrigin-RevId: 931228650

Author: ratgr

tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/BUILD

@@ -64,13 +64,13 @@ py_library(
     ],
     # strict_deps = True,
     deps = [
+        ":default_n_bit_quantize_configs",
+        ":default_n_bit_quantizers",
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/keras:compat",
         "//tensorflow_model_optimization/python/core/quantization/keras:quantize_config",
         "//tensorflow_model_optimization/python/core/quantization/keras:quantize_registry",
         "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
-        "//tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit:default_n_bit_quantize_configs",
-        "//tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit:default_n_bit_quantizers",
     ],
 )
 
@@ -88,6 +88,7 @@ py_test(
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/keras:compat",
         "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
+        "//tensorflow_model_optimization/python/core/quantization/keras:utils",
     ],
 )
 

tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantize_registry.py

@@ -20,7 +20,6 @@
 
 from typing import Any, Dict
 
-import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.keras.compat import keras
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_config
@@ -185,9 +184,15 @@ def __init__(self, disable_per_axis=False,
     self._num_bits_activation = num_bits_activation
     self._layer_quantize_map = {}
     for quantize_info in self._LAYER_QUANTIZE_INFO:
-      quantize_info.num_bits_weight = num_bits_weight
-      quantize_info.num_bits_activation = num_bits_activation
-      self._layer_quantize_map[quantize_info.layer_type] = quantize_info
+      new_quantize_info = _QuantizeInfo(
+          layer_type=quantize_info.layer_type,
+          weight_attrs=quantize_info.weight_attrs,
+          activation_attrs=quantize_info.activation_attrs,
+          quantize_output=quantize_info.quantize_output,
+          num_bits_weight=num_bits_weight,
+          num_bits_activation=num_bits_activation,
+      )
+      self._layer_quantize_map[new_quantize_info.layer_type] = new_quantize_info
 
     # Hack for `Activation` layer. That is the only layer with a separate
     # QuantizeConfig.
@@ -548,9 +553,13 @@ class DefaultNBitConvQuantizeConfig(DefaultNBitQuantizeConfig):
   def __init__(self, weight_attrs, activation_attrs,
                quantize_output, num_bits_weight: int = 8,
                num_bits_activation: int = 8):
-    super(DefaultNBitConvQuantizeConfig,
-          self).__init__(weight_attrs, activation_attrs,
-                         quantize_output, num_bits_weight, num_bits_activation)
+    super().__init__(
+        weight_attrs,
+        activation_attrs,
+        quantize_output,
+        num_bits_weight,
+        num_bits_activation,
+    )
     self._num_bits_weight = num_bits_weight
     self._num_bits_activation = num_bits_activation
     self.weight_quantizer = n_bit_quantizers.DefaultNBitConvWeightsQuantizer(
@@ -563,17 +572,22 @@ class DefaultNBitConvTransposeQuantizeConfig(
 
   def __init__(self, weight_attrs, activation_attrs, quantize_output,
                num_bits_weight: int = 8, num_bits_activation: int = 8):
-    super(DefaultNBitConvTransposeQuantizeConfig,
-          self).__init__(weight_attrs, activation_attrs, quantize_output,
-                         num_bits_weight, num_bits_activation)
+    super().__init__(
+        weight_attrs,
+        activation_attrs,
+        quantize_output,
+        num_bits_weight,
+        num_bits_activation,
+    )
     self._num_bits_weight = num_bits_weight
     self._num_bits_activation = num_bits_activation
 
     self.weight_quantizer = n_bit_quantizers.DefaultNBitConvTransposeWeightsQuantizer(
         num_bits_weight, num_bits_activation)
 
 
-def _types_dict():
+def types_dict():
+  """Returns a dictionary mapping type names to classes for deserialization."""
   return {
       'DefaultNBitQuantizeConfig':
           DefaultNBitQuantizeConfig,

tensorflow_model_optimization/python/core/quantization/keras/experimental/default_n_bit/default_n_bit_quantize_registry_test.py

@@ -78,7 +78,7 @@ class QuantizeRegistryTest(
     tf.test.TestCase, parameterized.TestCase, _TestHelper):
 
   def setUp(self):
-    super(QuantizeRegistryTest, self).setUp()
+    super().setUp()
     self.quantize_registry = n_bit_registry.DefaultNBitQuantizeRegistry(
         num_bits_weight=4, num_bits_activation=4)
 
@@ -114,7 +114,7 @@ class MinimalRNNCell(l.Layer):
       def __init__(self, units, **kwargs):
         self.units = units
         self.state_size = units
-        super(MinimalRNNCell, self).__init__(**kwargs)
+        super().__init__(**kwargs)
 
     self.assertFalse(
         self.quantize_registry.supports(l.RNN(cell=MinimalRNNCell(10))))
@@ -374,15 +374,16 @@ def testSerialization(self):
     quantize_config_from_config = deserialize_keras_object(
         serialized_quantize_config,
         module_objects=globals(),
-        custom_objects=n_bit_registry._types_dict())
+        custom_objects=n_bit_registry.types_dict(),
+    )
 
     self.assertEqual(quantize_config, quantize_config_from_config)
 
 
 class DefaultNBitQuantizeConfigRNNTest(tf.test.TestCase, _TestHelper):
 
   def setUp(self):
-    super(DefaultNBitQuantizeConfigRNNTest, self).setUp()
+    super().setUp()
 
     self.cell1 = l.LSTMCell(3)
     self.cell2 = l.GRUCell(2)
@@ -493,7 +494,8 @@ def testSerialization(self):
     quantize_config_from_config = deserialize_keras_object(
         serialized_quantize_config,
         module_objects=globals(),
-        custom_objects=n_bit_registry._types_dict())
+        custom_objects=n_bit_registry.types_dict(),
+    )
 
     self.assertEqual(self.quantize_config, quantize_config_from_config)
 

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

@@ -74,7 +74,7 @@ def quantize_scope(*args):
       'FixedQuantizeConfig': quantize_config_mod.FixedQuantizeConfig,
   }
   quantization_objects.update(default_8bit_quantize_registry._types_dict())  # pylint: disable=protected-access
-  quantization_objects.update(default_n_bit_quantize_registry._types_dict())  # pylint: disable=protected-access
+  quantization_objects.update(default_n_bit_quantize_registry.types_dict())
   quantization_objects.update(quantizers._types_dict())  # pylint: disable=protected-access
 
   return keras.utils.custom_object_scope(*(args + (quantization_objects,)))