Add DenseBatchNorm.* transforms to default_n_bit scheme

PiperOrigin-RevId: 428949606

Author: daverim

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

@@ -134,6 +134,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/experimental/default_n_bit/default_n_bit_quantize_layout_transform.py

@@ -104,6 +104,15 @@ def apply(self, model, layer_quantize_map):
         default_n_bit_transforms.LayerReluActivationQuantize(
             num_bits_weight=self._num_bits_weight,
             num_bits_activation=self._num_bits_activation),
+        default_n_bit_transforms.DenseBatchNormQuantize(
+            num_bits_weight=self._num_bits_weight,
+            num_bits_activation=self._num_bits_activation),
+        default_n_bit_transforms.DenseBatchNormReLUQuantize(
+            num_bits_weight=self._num_bits_weight,
+            num_bits_activation=self._num_bits_activation),
+        default_n_bit_transforms.DenseBatchNormActivationQuantize(
+            num_bits_weight=self._num_bits_weight,
+            num_bits_activation=self._num_bits_activation),
     ]
     return model_transformer.ModelTransformer(
         model, transforms,

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

@@ -263,6 +263,95 @@ def pattern(self):
         inputs=[Conv2DReshapeBatchNormQuantize.pattern(self)])
 
 
+class DenseBatchNormQuantize(transforms.Transform):
+  """Transform to be applied to "Dense"+ "BatchNorm" Graph.
+
+  This transform disables Quantization between Dense and BatchNorm
+  to ensure FQ does not get placed between them.
+  """
+
+  def __init__(self, num_bits_weight: int = 8, num_bits_activation: int = 8):
+    self._num_bits_weight = num_bits_weight
+    self._num_bits_activation = num_bits_activation
+
+  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'] = (
+        configs.DefaultNBitOutputQuantizeConfig(
+            num_bits_weight=self._num_bits_weight,
+            num_bits_activation=self._num_bits_activation))
+    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 {
+        'DefaultNBitOutputQuantizeConfig':
+            configs.DefaultNBitOutputQuantizeConfig,
+        'NoOpQuantizeConfig':
+            configs.NoOpQuantizeConfig,
+        'NoOpActivation': quantize_aware_activation.NoOpActivation
+    }
+
+
+class DenseBatchNormReLUQuantize(DenseBatchNormQuantize):
+  """Transform to be applied to "Dense"+ "BatchNorm" + "ReLU" Graph.
+
+  This transform disables Quantization between Dense, BatchNorm and ReLU
+  to ensure FQ does not get placed between them.
+  """
+
+  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'] = (
+        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):
+  """Transform to be applied to "Dense"+ "BatchNorm" + "ReLU" Graph.
+
+  This transform disables Quantization between Dense, BatchNorm and ReLU
+  to ensure FQ does not get placed between them.
+  """
+
+  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/experimental/default_n_bit/default_n_bit_transforms_test.py

@@ -29,11 +29,13 @@
 from tensorflow_model_optimization.python.core.quantization.keras.experimental.default_n_bit import default_n_bit_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,21 +75,26 @@ 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,
+      conv_or_dense_activation_class,
       normalization_quantize_config_class):
     model = self._get_model(layer_type,
                             squeeze_type,
@@ -107,7 +114,7 @@ def _test_conv_squeeze_bn_activation_transform(
       bn_layer = transformed_model.layers[2]
 
     self.assertIsInstance(
-        conv_layer.activation, conv_activation_class)
+        conv_layer.activation, conv_or_dense_activation_class)
     self.assertIsInstance(
         updated_metadata.get(bn_layer.name).get('quantize_config'),
         normalization_quantize_config_class)
@@ -123,13 +130,13 @@ 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_n_bit_transforms.Conv2DBatchNormQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
         normalization_quantize_config_class=
         n_bit_configs.DefaultNBitOutputQuantizeConfig)
 
@@ -140,14 +147,14 @@ 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_n_bit_transforms.Conv2DBatchNormReLUQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
         normalization_quantize_config_class=
         n_bit_configs.NoOpQuantizeConfig)
 
@@ -159,14 +166,14 @@ 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_n_bit_transforms.Conv2DBatchNormActivationQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
         normalization_quantize_config_class=
         n_bit_configs.NoOpQuantizeConfig)
 
@@ -178,14 +185,14 @@ 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_n_bit_transforms.Conv2DReshapeBatchNormQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
         normalization_quantize_config_class=
         n_bit_configs.DefaultNBitOutputQuantizeConfig)
 
@@ -197,14 +204,14 @@ 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_n_bit_transforms.Conv2DReshapeBatchNormReLUQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
         normalization_quantize_config_class=
         n_bit_configs.NoOpQuantizeConfig)
 
@@ -216,17 +223,64 @@ 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_n_bit_transforms.Conv2DReshapeBatchNormActivationQuantize,
-        conv_activation_class=quantize_aware_activation.NoOpActivation,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
         normalization_quantize_config_class=
         n_bit_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_n_bit_transforms.DenseBatchNormQuantize,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=n_bit_configs
+        .DefaultNBitOutputQuantizeConfig)
+
+  @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_n_bit_transforms.DenseBatchNormReLUQuantize,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=n_bit_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_n_bit_transforms
+        .DenseBatchNormActivationQuantize,
+        conv_or_dense_activation_class=quantize_aware_activation.NoOpActivation,
+        normalization_quantize_config_class=n_bit_configs
+        .NoOpQuantizeConfig)
+
   @parameterized.named_parameters(
       ('padding_valid', {'padding': 'valid'}),
       ('padding_same', {'padding': 'same'}),