Add DenseBatchNormReLU transform for non-folded case

PiperOrigin-RevId: 409027616

Author: dayeongl

tensorflow_model_optimization/python/core/quantization/keras/default_8bit/BUILD

@@ -122,6 +122,7 @@ py_strict_test(
         "//tensorflow_model_optimization/python/core/quantization/keras:quantizers",
         "//tensorflow_model_optimization/python/core/quantization/keras/graph_transformations:model_transformer",
         "//tensorflow_model_optimization/python/core/quantization/keras/layers:conv_batchnorm_test_utils",
+        "//tensorflow_model_optimization/python/core/quantization/keras/layers:dense_batchnorm_test_utils",
     ],
 )
 

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

@@ -65,6 +65,9 @@ def apply(self, model, layer_quantize_map):
         default_8bit_transforms.ConcatTransform4Inputs(),
         default_8bit_transforms.ConcatTransform3Inputs(),
         default_8bit_transforms.ConcatTransform(),
+        default_8bit_transforms.DenseBatchNormQuantize(),
+        default_8bit_transforms.DenseBatchNormReLUQuantize(),
+        default_8bit_transforms.DenseBatchNormActivationQuantize(),
         default_8bit_transforms.LayerReLUQuantize(),
         default_8bit_transforms.LayerReluActivationQuantize(),
     ]

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

@@ -221,6 +221,75 @@ def pattern(self):
         inputs=[Conv2DReshapeBatchNormQuantize.pattern(self)])
 
 
+class DenseBatchNormQuantize(transforms.Transform):
+  """Ensure FQ does not get placed between Dense and BatchNorm."""
+
+  def pattern(self):
+    return LayerPattern(
+        'BatchNormalization|SyncBatchNormalization',
+        inputs=[LayerPattern('Dense', config={'activation': 'linear'})])
+
+  def _replace(self, bn_layer_node, dense_layer_node):
+    if _has_custom_quantize_config(bn_layer_node, dense_layer_node):
+      return bn_layer_node
+
+    dense_layer_node.layer['config']['activation'] = (
+        keras.activations.serialize(quantize_aware_activation.NoOpActivation()))
+    bn_layer_node.metadata['quantize_config'] = (
+        default_8bit_quantize_configs.Default8BitOutputQuantizeConfig())
+
+    return bn_layer_node
+
+  def replacement(self, match_layer):
+    bn_layer_node = match_layer
+    dense_layer_node = match_layer.input_layers[0]
+
+    return self._replace(bn_layer_node, dense_layer_node)
+
+  def custom_objects(self):
+    return {
+        'NoOpQuantizeConfig': default_8bit_quantize_configs.NoOpQuantizeConfig,
+        'NoOpActivation': quantize_aware_activation.NoOpActivation
+    }
+
+
+class DenseBatchNormReLUQuantize(DenseBatchNormQuantize):
+  """Ensure FQ does not get placed between Dense, BatchNorm and ReLU."""
+
+  def pattern(self):
+    return LayerPattern(
+        'ReLU', inputs=[super(DenseBatchNormReLUQuantize, self).pattern()])
+
+  def _replace(self, relu_layer_node, bn_layer_node, dense_layer_node):
+    if _has_custom_quantize_config(relu_layer_node, bn_layer_node,
+                                   dense_layer_node):
+      return relu_layer_node
+
+    dense_layer_node.layer['config']['activation'] = (
+        keras.activations.serialize(quantize_aware_activation.NoOpActivation()))
+    bn_layer_node.metadata['quantize_config'] = (
+        default_8bit_quantize_configs.NoOpQuantizeConfig())
+
+    return relu_layer_node
+
+  def replacement(self, match_layer):
+    relu_layer_node = match_layer
+    bn_layer_node = relu_layer_node.input_layers[0]
+    dense_layer_node = bn_layer_node.input_layers[0]
+
+    return self._replace(relu_layer_node, bn_layer_node, dense_layer_node)
+
+
+class DenseBatchNormActivationQuantize(DenseBatchNormReLUQuantize):
+  """Ensure FQ does not get placed between Dense, BatchNorm and ReLU."""
+
+  def pattern(self):
+    return LayerPattern(
+        'Activation',
+        config={'activation': 'relu'},
+        inputs=[DenseBatchNormQuantize.pattern(self)])
+
+
 class SeparableConv1DQuantize(transforms.Transform):
   """Add QAT support for Keras SeparableConv1D layer.
 

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

@@ -29,11 +29,13 @@
 from tensorflow_model_optimization.python.core.quantization.keras.default_8bit import default_8bit_transforms
 from tensorflow_model_optimization.python.core.quantization.keras.graph_transformations import model_transformer
 from tensorflow_model_optimization.python.core.quantization.keras.layers import conv_batchnorm_test_utils
+from tensorflow_model_optimization.python.core.quantization.keras.layers import dense_batchnorm_test_utils
 
 ModelTransformer = model_transformer.ModelTransformer
 
 Conv2DModel = conv_batchnorm_test_utils.Conv2DModel
 DepthwiseConv2DModel = conv_batchnorm_test_utils.DepthwiseConv2DModel
+DenseModel = dense_batchnorm_test_utils.DenseModel
 
 keras = tf.keras
 
@@ -73,22 +75,27 @@ def _get_model(
           post_bn_activation=activation,
           squeeze_type=squeeze_type,
           normalization_type=normalization_type)
+    elif layer_type == 'Dense':
+      return DenseModel.get_nonfolded_batchnorm_model(
+          post_bn_activation=activation, normalization_type=normalization_type)
 
   def _get_input_shape(self, layer_type):
     if layer_type == 'Conv2D':
       return Conv2DModel.get_batched_input_shape()
     elif layer_type == 'DepthwiseConv2D':
       return DepthwiseConv2DModel.get_batched_input_shape()
+    elif layer_type == 'Dense':
+      return DenseModel.get_batched_input_shape()
 
-  def _test_conv_squeeze_bn_activation_transform(
+  def _test_conv_squeeze_or_dense_bn_activation_transform(
       self,
       layer_type,
-      squeeze_type,
       normalization_type,
       activation_type,
       transform_class,
-      conv_activation_class,
-      normalization_quantize_config_class):
+      conv_or_dense_activation_class,
+      normalization_quantize_config_class,
+      squeeze_type=None):
     model = self._get_model(layer_type,
                             squeeze_type,
                             normalization_type,
@@ -100,14 +107,14 @@ def _test_conv_squeeze_bn_activation_transform(
         [transform_class()],
     ).transform()
 
-    conv_layer = transformed_model.layers[1]
+    conv_or_dense_layer = transformed_model.layers[1]
     if squeeze_type == 'sepconv1d_squeeze':
       bn_layer = transformed_model.layers[3]
     else:
       bn_layer = transformed_model.layers[2]
 
-    self.assertIsInstance(
-        conv_layer.activation, conv_activation_class)
+    self.assertIsInstance(conv_or_dense_layer.activation,
+                          conv_or_dense_activation_class)
     self.assertIsInstance(
         updated_metadata.get(bn_layer.name).get('quantize_config'),
         normalization_quantize_config_class)
@@ -123,15 +130,15 @@ def _test_conv_squeeze_bn_activation_transform(
       ('DepthwiseConv2D', 'SyncBatchNormalization'),
   )
   def testConv2DBatchNormQuantize(self, layer_type, normalization_type):
-    self._test_conv_squeeze_bn_activation_transform(
+    self._test_conv_squeeze_or_dense_bn_activation_transform(
         layer_type=layer_type,
         squeeze_type=None,
         normalization_type=normalization_type,
         activation_type=None,
         transform_class=default_8bit_transforms.Conv2DBatchNormQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
-        normalization_quantize_config_class=
-        default_8bit_quantize_configs.Default8BitOutputQuantizeConfig)
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=default_8bit_quantize_configs
+        .Default8BitOutputQuantizeConfig)
 
   @parameterized.parameters(
       ('Conv2D', 'BatchNormalization'),
@@ -140,16 +147,15 @@ def testConv2DBatchNormQuantize(self, layer_type, normalization_type):
       ('DepthwiseConv2D', 'SyncBatchNormalization'),
   )
   def testConv2DBatchNormReLUQuantize(self, layer_type, normalization_type):
-    self._test_conv_squeeze_bn_activation_transform(
+    self._test_conv_squeeze_or_dense_bn_activation_transform(
         layer_type=layer_type,
         squeeze_type=None,
         normalization_type=normalization_type,
         activation_type='relu',
-        transform_class=
-        default_8bit_transforms.Conv2DBatchNormReLUQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
-        normalization_quantize_config_class=
-        default_8bit_quantize_configs.NoOpQuantizeConfig)
+        transform_class=default_8bit_transforms.Conv2DBatchNormReLUQuantize,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=default_8bit_quantize_configs
+        .NoOpQuantizeConfig)
 
   @parameterized.parameters(
       ('Conv2D', 'BatchNormalization'),
@@ -159,16 +165,16 @@ def testConv2DBatchNormReLUQuantize(self, layer_type, normalization_type):
   )
   def testConv2DBatchNormActivationQuantize(
       self, layer_type, normalization_type):
-    self._test_conv_squeeze_bn_activation_transform(
+    self._test_conv_squeeze_or_dense_bn_activation_transform(
         layer_type=layer_type,
         squeeze_type=None,
         normalization_type=normalization_type,
         activation_type='act_relu',
-        transform_class=
-        default_8bit_transforms.Conv2DBatchNormActivationQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
-        normalization_quantize_config_class=
-        default_8bit_quantize_configs.NoOpQuantizeConfig)
+        transform_class=default_8bit_transforms
+        .Conv2DBatchNormActivationQuantize,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=default_8bit_quantize_configs
+        .NoOpQuantizeConfig)
 
   @parameterized.parameters(
       ('Conv2D', 'BatchNormalization'),
@@ -178,16 +184,15 @@ def testConv2DBatchNormActivationQuantize(
   )
   def testConv2DReshapeBatchNormQuantize(
       self, layer_type, normalization_type):
-    self._test_conv_squeeze_bn_activation_transform(
+    self._test_conv_squeeze_or_dense_bn_activation_transform(
         layer_type=layer_type,
         squeeze_type='sepconv1d_squeeze',
         normalization_type=normalization_type,
         activation_type=False,
-        transform_class=
-        default_8bit_transforms.Conv2DReshapeBatchNormQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
-        normalization_quantize_config_class=
-        default_8bit_quantize_configs.Default8BitOutputQuantizeConfig)
+        transform_class=default_8bit_transforms.Conv2DReshapeBatchNormQuantize,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=default_8bit_quantize_configs
+        .Default8BitOutputQuantizeConfig)
 
   @parameterized.parameters(
       ('Conv2D', 'BatchNormalization'),
@@ -197,16 +202,16 @@ def testConv2DReshapeBatchNormQuantize(
   )
   def testConv2DReshapeBatchNormReLUQuantize(
       self, layer_type, normalization_type):
-    self._test_conv_squeeze_bn_activation_transform(
+    self._test_conv_squeeze_or_dense_bn_activation_transform(
         layer_type=layer_type,
         squeeze_type='sepconv1d_squeeze',
         normalization_type=normalization_type,
         activation_type='relu',
-        transform_class=
-        default_8bit_transforms.Conv2DReshapeBatchNormReLUQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
-        normalization_quantize_config_class=
-        default_8bit_quantize_configs.NoOpQuantizeConfig)
+        transform_class=default_8bit_transforms
+        .Conv2DReshapeBatchNormReLUQuantize,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=default_8bit_quantize_configs
+        .NoOpQuantizeConfig)
 
   @parameterized.parameters(
       ('Conv2D', 'BatchNormalization'),
@@ -216,16 +221,63 @@ def testConv2DReshapeBatchNormReLUQuantize(
   )
   def testConv2DReshapeBatchNormActivationQuantize(
       self, layer_type, normalization_type):
-    self._test_conv_squeeze_bn_activation_transform(
+    self._test_conv_squeeze_or_dense_bn_activation_transform(
         layer_type=layer_type,
         squeeze_type='sepconv1d_squeeze',
         normalization_type=normalization_type,
         activation_type='act_relu',
-        transform_class=
-        default_8bit_transforms.Conv2DReshapeBatchNormActivationQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
-        normalization_quantize_config_class=
-        default_8bit_quantize_configs.NoOpQuantizeConfig)
+        transform_class=default_8bit_transforms
+        .Conv2DReshapeBatchNormActivationQuantize,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=default_8bit_quantize_configs
+        .NoOpQuantizeConfig)
+
+  @parameterized.parameters(
+      ('Dense', 'BatchNormalization'),
+      ('Dense', 'SyncBatchNormalization'),
+  )
+  def testDenseBatchNormQuantize(self, layer_type, normalization_type):
+    self._test_conv_squeeze_or_dense_bn_activation_transform(
+        layer_type=layer_type,
+        squeeze_type=None,
+        normalization_type=normalization_type,
+        activation_type=None,
+        transform_class=default_8bit_transforms.DenseBatchNormQuantize,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=default_8bit_quantize_configs
+        .Default8BitOutputQuantizeConfig)
+
+  @parameterized.parameters(
+      ('Dense', 'BatchNormalization'),
+      ('Dense', 'SyncBatchNormalization'),
+  )
+  def testDenseBatchNormReLUQuantize(self, layer_type, normalization_type):
+    self._test_conv_squeeze_or_dense_bn_activation_transform(
+        layer_type=layer_type,
+        squeeze_type=None,
+        normalization_type=normalization_type,
+        activation_type='relu',
+        transform_class=default_8bit_transforms.DenseBatchNormReLUQuantize,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=default_8bit_quantize_configs
+        .NoOpQuantizeConfig)
+
+  @parameterized.parameters(
+      ('Dense', 'BatchNormalization'),
+      ('Dense', 'SyncBatchNormalization'),
+  )
+  def testDenseBatchNormActivationQuantize(self, layer_type,
+                                           normalization_type):
+    self._test_conv_squeeze_or_dense_bn_activation_transform(
+        layer_type=layer_type,
+        squeeze_type=None,
+        normalization_type=normalization_type,
+        activation_type='act_relu',
+        transform_class=default_8bit_transforms
+        .DenseBatchNormActivationQuantize,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=default_8bit_quantize_configs
+        .NoOpQuantizeConfig)
 
   @parameterized.named_parameters(
       ('padding_valid', {'padding': 'valid'}),

tensorflow_model_optimization/python/core/quantization/keras/layers/BUILD

@@ -23,3 +23,12 @@ py_strict_library(
         # tensorflow dep1,
     ],
 )
+
+py_strict_library(
+    name = "dense_batchnorm_test_utils",
+    srcs = ["dense_batchnorm_test_utils.py"],
+    srcs_version = "PY3",
+    deps = [
+        # tensorflow dep1,
+    ],
+)