Update Concat transforms to handle QuantizeLayer inputs by removing quantizer. Allow quantize layer to take None quantizers.

PiperOrigin-RevId: 406065456

Author: daverim

tensorflow_model_optimization/python/core/quantization/keras/default_8bit/default_8bit_transforms.py

@@ -504,6 +504,8 @@ def pattern(self):
         'Concatenate', inputs=[LayerPattern('.*'), LayerPattern('.*')])
 
   def _get_layer_type(self, layer_class_name):
+    if layer_class_name == 'QuantizeLayer':
+      return quantize_layer.QuantizeLayer
     keras_layers = inspect.getmembers(tf.keras.layers, inspect.isclass)
     for layer_name, layer_type in keras_layers:
       if layer_name == layer_class_name:
@@ -537,6 +539,10 @@ def replacement(self, match_layer):
               default_8bit_quantize_configs.NoOpQuantizeConfig())
           continue
 
+        if layer_class == quantize_layer.QuantizeLayer:
+          feed_layer_node.metadata['quantizer'] = None
+          continue
+
         if not default_registry._is_supported_layer(layer_class):
           # Feeding layer is not supported by registry
           return match_layer

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

@@ -375,6 +375,18 @@ def _quantize(layer):  # pylint: disable=missing-docstring
       # It supports for custom QuantizeWrapper.
       return layer
 
+    # layer is a QuantizeLayer, possibly rebuild
+    # layer with modified config from parameters stored in the map.
+    if isinstance(layer, quantize_layer.QuantizeLayer):
+      # If there is more than one usage of the input, even if all are concat,
+      # we need to quantize.
+      if len(layer._outbound_nodes) > 1:  # pylint: disable=protected-access
+        return layer
+      layer_config = layer.get_config()
+      for key, value in layer_quantize_map[layer.name].items():
+        layer_config[key] = value
+      return quantize_layer.QuantizeLayer.from_config(layer_config)
+
     if layer.name in requires_output_quantize:
       if not quantize_registry.supports(layer):
         return layer

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

@@ -38,20 +38,23 @@ def __init__(self, quantizer, **kwargs):
     """Create a QuantizeLayer.
 
     Args:
-      quantizer: `Quantizer` used to quantize tensors.
+      quantizer: `Quantizer` used to quantize tensors. quantizer=None
+        means no quantization of the input layer.
       **kwargs: Additional keyword arguments to be passed to the keras layer.
     """
     super(QuantizeLayer, self).__init__(**kwargs)
 
-    if quantizer is None or not isinstance(quantizer, quantizers.Quantizer):
-      raise ValueError('quantizer should not be None, and should be an instance'
+    if quantizer is not None and not isinstance(quantizer,
+                                                quantizers.Quantizer):
+      raise ValueError('quantizer should be an instance'
                        'of `tfmot.quantization.keras.quantizers.Quantizer`.')
 
     self.quantizer = quantizer
 
   def build(self, input_shape):
-    self.quantizer_vars = self.quantizer.build(
-        input_shape, self.name, self)
+    if self.quantizer:
+      self.quantizer_vars = self.quantizer.build(
+          input_shape, self.name, self)
 
     self.optimizer_step = self.add_weight(
         'optimizer_step',
@@ -60,6 +63,9 @@ def build(self, input_shape):
         trainable=False)
 
   def call(self, inputs, training=None):
+    if not self.quantizer:
+      return inputs
+
     if training is None:
       training = tf.keras.backend.learning_phase()
 

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

@@ -69,6 +69,22 @@ def testSerializationQuantizeLayer(self):
 
     self.assertEqual(layer_from_config.get_config(), layer.get_config())
 
+  def testNoQuantizeLayer(self):
+    layer = QuantizeLayer(quantizer=None, input_shape=(4,))
+    model = tf.keras.Sequential([layer])
+    x = np.random.rand(1, 4)
+    self.assertAllClose(x, model.predict(x))
+
+    custom_objects = {
+        'QuantizeLayer': QuantizeLayer,
+    }
+
+    serialized_layer = serialize_layer(layer)
+    with tf.keras.utils.custom_object_scope(custom_objects):
+      layer_from_config = deserialize_layer(serialized_layer)
+
+    self.assertEqual(layer_from_config.get_config(), layer.get_config())
+
 
 if __name__ == '__main__':
   tf.test.main()