Add QuantizeWrapperV2 that preserve weights order, and make it as a default for quantize_apply.

PiperOrigin-RevId: 396525480

Author: Xhark

tensorflow_model_optimization/python/core/api/quantization/keras/__init__.py

@@ -34,7 +34,7 @@
 # handle custom Keras layers.
 from tensorflow_model_optimization.python.core.quantization.keras.quantize_config import QuantizeConfig
 from tensorflow_model_optimization.python.core.quantization.keras.quantize_wrapper import QuantizeWrapper
-
+from tensorflow_model_optimization.python.core.quantization.keras.quantize_wrapper import QuantizeWrapperV2
 # Deserialize quantized model for Keras h5 format.
 from tensorflow_model_optimization.python.core.quantization.keras.quantize import quantize_scope
 

tensorflow_model_optimization/python/core/quantization/keras/BUILD

@@ -269,6 +269,7 @@ py_strict_test(
         ":quantize_config",
         ":quantize_layer",
         ":quantize_wrapper",
+        ":quantizers",
         # numpy dep1,
         # tensorflow dep1,
         "//tensorflow_model_optimization/python/core/keras:test_utils",

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

@@ -67,6 +67,7 @@ def quantize_scope(*args):
           quantize_aware_activation.QuantizeAwareActivation,
       'NoOpActivation': quantize_aware_activation.NoOpActivation,
       'QuantizeWrapper': quantize_wrapper.QuantizeWrapper,
+      'QuantizeWrapperV2': quantize_wrapper.QuantizeWrapperV2,
       'QuantizeLayer': quantize_layer.QuantizeLayer,
       'OutputOnlyConfig': quantize_config_mod.OutputOnlyConfig,
   }
@@ -401,7 +402,8 @@ def _quantize(layer):  # pylint: disable=missing-docstring
     # `QuantizeAnnotate`. This should generally be fine, but occasionally
     # `QuantizeAnnotate` wrapper may contain `batch_input_shape` like params.
     # TODO(pulkitb): Ensure this does not affect model cloning.
-    return quantize_wrapper.QuantizeWrapper(layer, quantize_config)
+    return quantize_wrapper.QuantizeWrapperV2(
+        layer, quantize_config)
 
   # 1. Create a copy of the model with the same weights. This ensures
   # modifications don't affect the original model, or its weights.

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

@@ -27,6 +27,7 @@
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_config as quantize_config_mod
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_layer
 from tensorflow_model_optimization.python.core.quantization.keras import quantize_wrapper as quantize_wrapper_mod
+from tensorflow_model_optimization.python.core.quantization.keras import quantizers
 from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_quantize_registry
 
 quantize_annotate_layer = quantize.quantize_annotate_layer
@@ -513,6 +514,68 @@ def testQuantizeApply_RunsWhenNestedModelNotAnnotated(self):
 
     quantize_apply(annotated_model)
 
+  class CustomConvLayer(tf.keras.layers.Layer):
+
+    def __init__(self, name=None, **kwargs):
+      super().__init__(name=name, **kwargs)
+      self.conv1 = tf.keras.layers.Conv2D(2, 2)
+
+    def build(self, input_shape):
+      self.conv1.build(input_shape)
+
+    def call(self, inputs):
+      return self.conv1(inputs)
+
+    def get_config(self):
+      return {'name': self.name}
+
+  class CustomConvQuantizeConfig(quantize_config_mod.QuantizeConfig):
+
+    def get_weights_and_quantizers(self, layer):
+      return [(layer.conv1.kernel, quantizers.LastValueQuantizer(
+          num_bits=8, symmetric=True, narrow_range=False, per_axis=False)),]
+
+    def get_activations_and_quantizers(self, layer):
+      return []
+
+    def set_quantize_weights(self, layer, quantize_weights):
+      # layer.conv1._kernel_bak = layer.conv1.kernel
+      layer.conv1.kernel = quantize_weights[0]
+
+    def set_quantize_activations(self, layer, quantize_activations):
+      pass
+
+    def get_output_quantizers(self, layer):
+      return []
+
+    def get_config(self):
+      return {}
+
+  def testQuantizeApply_KeepTrainableWeightOrder(self):
+    layer = self.CustomConvLayer(input_shape=(28, 28, 3))
+    model = keras.Sequential([layer])
+
+    def apply_quantization_to_dense(layer):
+      if isinstance(layer, self.CustomConvLayer):
+        return quantize_annotate_layer(
+            layer, quantize_config=self.CustomConvQuantizeConfig())
+      return layer
+
+    annotated_model = tf.keras.models.clone_model(
+        model,
+        clone_function=apply_quantization_to_dense,
+    )
+
+    with quantize.quantize_scope({
+        'CustomConvQuantizeConfig': self.CustomConvQuantizeConfig,
+        'CustomConvLayer': self.CustomConvLayer
+    }):
+      quant_aware_model = quantize_apply(annotated_model)
+
+    self._assert_weights_different_objects(
+        model.trainable_weights, quant_aware_model.trainable_weights)
+    self._assert_weights_equal_value(
+        model.trainable_weights, quant_aware_model.trainable_weights)
 
 if __name__ == '__main__':
   tf.test.main()

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

@@ -228,3 +228,19 @@ def updates(self):
   @property
   def losses(self):
     return self.layer.losses + self._losses
+
+
+# TODO(b/199809494): Update guide document to use QuantizeWrapperV2.
+# Do not override this class method to quantize wrapper directly.
+# It breaks existing h5 models that uses QuantizeWrapper class.
+class QuantizeWrapperV2(QuantizeWrapper):
+
+  def build(self, input_shape):
+    self._trainable_weights.extend(self.layer.trainable_weights)
+    super(QuantizeWrapperV2, self).build(input_shape)
+
+  @property
+  def trainable_weights(self):
+    # Change the order to keep the weight order after applying QAT.
+    return self._dedup_weights(
+        self._trainable_weights + self.layer.trainable_weights)