Support for SeparableConv2D for QAT

SeparableConv2D is a combination of DepthwiseConv2D
and Conv2D. QAT needs to break the layer up into this
combination so that the rest of the infra can then apply
QAT to the resulting layers.

Not possible to implement support for SeparableConv directly
since it does not provide hooks to alter internal graph
construction.

PiperOrigin-RevId: 321030603

Author: nutsiepully

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

@@ -228,6 +228,105 @@ def pattern(self):
         inputs=[Conv2DBatchNormQuantize.pattern(self)])
 
 
+class SeparableConvQuantize(transforms.Transform):
+  """Break SeparableConv into a DepthwiseConv and Conv layer.
+
+  SeparableConv is a composition of a DepthwiseConv and a Conv layer. For the
+  purpose of quantization, a FQ operation needs to be placed between the output
+  of DepthwiseConv and the following Conv.
+
+  This is needed since there is a dynamic tensor in between the two layers, and
+  it's range information needs to be captured by the FakeQuant op to ensure
+  full int8 quantization of the layers is possible.
+
+  Splitting the layer into 2 ensures that each individual layer is handled
+  correctly with respect to quantization.
+  """
+
+  def pattern(self):
+    return LayerPattern('SeparableConv2D')
+
+  @staticmethod
+  def _get_quantize_config(layer_node):
+    return layer_node.metadata.get('quantize_config')
+
+  def _has_custom_quantize_config(self, *layer_nodes):
+    for layer_node in layer_nodes:
+      if self._get_quantize_config(layer_node) is not None:
+        return True
+    return False
+
+  def replacement(self, match_layer):
+    if self._has_custom_quantize_config(match_layer):
+      return match_layer
+
+    sepconv_layer = match_layer.layer
+    sepconv_weights = list(match_layer.weights.values())
+
+    # TODO(pulkitb): SeparableConv has kwargs other than constructor args which
+    # need to be handled.
+    # Applicable to both layers: trainable, dtype, name
+    # Applicable to dconv: input_dim, input_shape, batch_input_shape, batch_size
+    # Needs special handling: weights
+    # Unknown: dynamic, autocast
+
+    dconv_layer = tf.keras.layers.DepthwiseConv2D(
+        kernel_size=sepconv_layer['config']['kernel_size'],
+        strides=sepconv_layer['config']['strides'],
+        padding=sepconv_layer['config']['padding'],
+        depth_multiplier=sepconv_layer['config']['depth_multiplier'],
+        data_format=sepconv_layer['config']['data_format'],
+        dilation_rate=sepconv_layer['config']['dilation_rate'],
+        activation=None,
+        use_bias=False,
+        depthwise_initializer=sepconv_layer['config']['depthwise_initializer'],
+        depthwise_regularizer=sepconv_layer['config']['depthwise_regularizer'],
+        depthwise_constraint=sepconv_layer['config']['depthwise_constraint'],
+        trainable=sepconv_layer['config']['trainable']
+    )
+    dconv_weights = collections.OrderedDict()
+    dconv_weights['depthwise_kernel:0'] = sepconv_weights[0]
+    dconv_layer_config = keras.layers.serialize(dconv_layer)
+    dconv_layer_config['name'] = dconv_layer.name
+    # Needed to ensure these new layers are considered for quantization.
+    dconv_metadata = {'quantize_config': None}
+
+    conv_layer = tf.keras.layers.Conv2D(
+        filters=sepconv_layer['config']['filters'],
+        kernel_size=(1, 1),  # (1,) * rank
+        strides=(1, 1),
+        padding='valid',
+        data_format=sepconv_layer['config']['data_format'],
+        dilation_rate=sepconv_layer['config']['dilation_rate'],
+        groups=1,
+        activation=sepconv_layer['config']['activation'],
+        use_bias=sepconv_layer['config']['use_bias'],
+        kernel_initializer=sepconv_layer['config']['pointwise_initializer'],
+        bias_initializer=sepconv_layer['config']['bias_initializer'],
+        kernel_regularizer=sepconv_layer['config']['pointwise_regularizer'],
+        bias_regularizer=sepconv_layer['config']['bias_regularizer'],
+        activity_regularizer=sepconv_layer['config']['activity_regularizer'],
+        kernel_constraint=sepconv_layer['config']['pointwise_constraint'],
+        bias_constraint=sepconv_layer['config']['bias_constraint'],
+        trainable=sepconv_layer['config']['trainable']
+    )
+    conv_weights = collections.OrderedDict()
+    conv_weights['kernel:0'] = sepconv_weights[1]
+    conv_weights['bias:0'] = sepconv_weights[2]
+    conv_layer_config = keras.layers.serialize(conv_layer)
+    conv_layer_config['name'] = conv_layer.name
+    # Needed to ensure these new layers are considered for quantization.
+    conv_metadata = {'quantize_config': None}
+
+    dconv_layer_node = LayerNode(
+        dconv_layer_config, weights=dconv_weights, metadata=dconv_metadata)
+    return LayerNode(
+        conv_layer_config,
+        weights=conv_weights,
+        input_layers=[dconv_layer_node],
+        metadata=conv_metadata)
+
+
 class AddReLUQuantize(transforms.Transform):
   """Ensure FQ does not get placed between Add and ReLU."""
 

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

@@ -221,6 +221,63 @@ def testConv2DBatchNormReLUQuantize(
     self.assertAllClose(
         transformed_model.predict(inputs), model.predict(inputs))
 
+  @parameterized.named_parameters(
+      ('padding_valid', {'padding': 'valid'}),
+      ('padding_same', {'padding': 'same'}),
+      ('padding_same_dilation_2', {'padding': 'same', 'dilation_rate': 2}),
+      ('strides', {'strides': 2}),
+      ('dilation_rate', {'dilation_rate': 2}),
+      ('depth_multiplier', {'depth_multiplier': 2}),
+      ('regularizer', {
+          'depthwise_regularizer': 'l2',
+          'pointwise_regularizer': 'l2',
+          'bias_regularizer': 'l2',
+          'activity_regularizer': 'l2'}),
+      ('constraint', {
+          'depthwise_constraint': tf.keras.constraints.max_norm(2.),
+          'pointwise_constraint': tf.keras.constraints.min_max_norm(0., 2.),
+          'bias_constraint': tf.keras.constraints.unit_norm()})
+  )
+  def testSeparableConvQuantize_(self, kwargs):
+    kwargs['filters'] = 2
+    kwargs['kernel_size'] = 3
+    num_samples = 2
+    stack_size = 3
+    num_row = 7
+    num_col = 6
+
+    sepconv_model = tf.keras.Sequential([
+        tf.keras.Input(
+            shape=(num_row, num_col, stack_size), batch_size=num_samples),
+        tf.keras.layers.SeparableConv2D(**kwargs)])
+
+    transformed_model, updated_metadata = ModelTransformer(
+        sepconv_model,
+        [default_8bit_transforms.SeparableConvQuantize()],
+    ).transform()
+
+    self.assertContainsSubset(
+        updated_metadata.keys(), {'depthwise_conv2d', 'conv2d'})
+    # Transformed model should have the same output shape
+    self.assertEqual(sepconv_model.output_shape, transformed_model.output_shape)
+
+    x = np.random.rand(*sepconv_model.input_shape)
+    y = np.random.rand(*sepconv_model.output_shape)
+
+    # Ensure model is equivalent, and forward pass results are the same.
+    self.assertAllClose(sepconv_model.predict(x), transformed_model.predict(x))
+
+    # Ensure model is equivalent, and training results are the same.
+    sepconv_model.compile(loss='categorical_crossentropy', optimizer='sgd')
+    sepconv_model.fit(x, y, epochs=100)
+    transformed_model.compile(loss='categorical_crossentropy', optimizer='sgd')
+    transformed_model.fit(x, y, epochs=100)
+
+    # Over a long training cycle with constraints and regularizers, the model
+    # can build very minute differences. Hence reducing tol to 1e-5.
+    self.assertAllClose(sepconv_model.predict(x), transformed_model.predict(x),
+                        atol=1e-5, rtol=1e-5)
+
   @parameterized.parameters(
       ('relu', default_8bit_transforms.AddReLUQuantize),
       ('act_relu', default_8bit_transforms.AddActivationQuantize),

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/transforms.py

@@ -41,7 +41,7 @@ class LayerPattern(object):
     pattern = LayerPattern('Concat', {}, [
         LayerPattern('Conv2D', {}, []),
         LayerPattern('Conv2D', {}, [])
-    )
+    ])
   """
 
   def __init__(self, class_name, config=None, inputs=None):