Add support for Relu folding in Conv/DConv/Dense

This makes ReLU6 support better since users typically
use 'relu' directly.

PiperOrigin-RevId: 364423101

Author: nutsiepully

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

@@ -65,8 +65,8 @@ def apply(self, model, layer_quantize_map):
         default_8bit_transforms.ConcatTransform4Inputs(),
         default_8bit_transforms.ConcatTransform3Inputs(),
         default_8bit_transforms.ConcatTransform(),
-        default_8bit_transforms.AddReLUQuantize(),
-        default_8bit_transforms.AddActivationQuantize(),
+        default_8bit_transforms.LayerReLUQuantize(),
+        default_8bit_transforms.LayerReluActivationQuantize(),
     ]
     return model_transformer.ModelTransformer(
         model, transforms,

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

@@ -497,11 +497,12 @@ def replacement(self, match_layer):
         metadata=conv_metadata)
 
 
-class AddReLUQuantize(transforms.Transform):
+class LayerReLUQuantize(transforms.Transform):
   """Ensure FQ does not get placed between Add and ReLU."""
 
   def pattern(self):
-    return LayerPattern('ReLU', inputs=[LayerPattern('Add')])
+    return LayerPattern(
+        'ReLU', inputs=[LayerPattern('Add|Conv2D|DepthwiseConv2D|Dense')])
 
   def replacement(self, match_layer):
     relu_layer_node = match_layer
@@ -518,14 +519,14 @@ def custom_objects(self):
     }
 
 
-class AddActivationQuantize(AddReLUQuantize):
+class LayerReluActivationQuantize(LayerReLUQuantize):
   """Ensure FQ does not get placed between Add and ReLU."""
 
   def pattern(self):
     return LayerPattern(
         'Activation',
         config={'activation': 'relu'},
-        inputs=[LayerPattern('Add')])
+        inputs=[LayerPattern('Add|Conv2D|DepthwiseConv2D|Dense')])
 
 
 class InputLayerQuantize(transforms.Transform):

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

@@ -344,8 +344,8 @@ def testSeparableConvQuantize_(self, kwargs):
   # Conv2DReshapeBatchNormActivationQuantize
 
   @parameterized.parameters(
-      ('relu', default_8bit_transforms.AddReLUQuantize),
-      ('act_relu', default_8bit_transforms.AddActivationQuantize),
+      ('relu', default_8bit_transforms.LayerReLUQuantize),
+      ('act_relu', default_8bit_transforms.LayerReluActivationQuantize),
   )
   def testAddReLUQuantize(self, activation_type, transform_type):
     add = keras.layers.Add()
@@ -370,6 +370,33 @@ def testAddReLUQuantize(self, activation_type, transform_type):
         updated_metadata.get(add_layer.name).get('quantize_config'),
         default_8bit_quantize_configs.NoOpQuantizeConfig)
 
+  @parameterized.parameters(
+      ('relu', default_8bit_transforms.LayerReLUQuantize),
+      ('act_relu', default_8bit_transforms.LayerReluActivationQuantize))
+  def testLayerReLUQuantize(self, activation_type, transform_type):
+    # TODO(tfmot): Add tests for DepthConv and Dense
+    input_shape = (1, 3, 3, 3)
+    conv_layer = tf.keras.layers.Conv2D(5, 2, input_shape=input_shape)
+    if activation_type == 'relu':
+      act_layer = keras.layers.ReLU(6.0)
+    elif activation_type == 'act_relu':
+      act_layer = keras.layers.Activation('relu')
+
+    model = tf.keras.Sequential([conv_layer, act_layer])
+
+    transformed_model, updated_metadata = ModelTransformer(
+        model,
+        [transform_type()],
+    ).transform()
+
+    self.assertIsInstance(
+        updated_metadata.get(model.layers[0].name).get('quantize_config'),
+        default_8bit_quantize_configs.NoOpQuantizeConfig)
+
+    inputs = np.random.standard_normal(input_shape)
+    self.assertAllClose(
+        transformed_model.predict(inputs), model.predict(inputs))
+
   def testAddsQuantizeLayerAfterInputLayer(self):
     inp1 = keras.layers.Input((3,))
     inp2 = keras.layers.Input((3,))