Add Conv2DTranspose supports.

PiperOrigin-RevId: 353149071

Author: Xhark

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

@@ -92,7 +92,6 @@ class Default8BitQuantizeRegistry(
       # layers.DepthwiseConv2D is supported and handled in code below.
 
       # _QuantizeInfo(layers.Conv3D, ['kernel'], ['activation']),
-      # _QuantizeInfo(layers.Conv2DTranspose, ['kernel'], ['activation']),
       # _QuantizeInfo(layers.Conv3DTranspose, ['kernel'], ['activation']),
       _no_quantize(layers.Cropping1D),
       _no_quantize(layers.Cropping2D),
@@ -198,6 +197,9 @@ def __init__(self):
     self._layer_quantize_map[
         layers.DepthwiseConv2D] = Default8BitConvQuantizeConfig(
             ['depthwise_kernel'], ['activation'], False)
+    self._layer_quantize_map[layers.Conv2DTranspose] = \
+        Default8BitConvTransposeQuantizeConfig(
+            ['kernel'], ['activation'], False)
 
   def _is_supported_layer(self, layer_class):
     return layer_class in self._layer_quantize_map
@@ -509,6 +511,17 @@ def __init__(self, weight_attrs, activation_attrs, quantize_output):
     )
 
 
+class Default8BitConvTransposeQuantizeConfig(Default8BitQuantizeConfig):
+  """QuantizeConfig for Conv2DTranspose layers."""
+
+  def __init__(self, weight_attrs, activation_attrs, quantize_output):
+    super(Default8BitConvTransposeQuantizeConfig,
+          self).__init__(weight_attrs, activation_attrs, quantize_output)
+
+    self.weight_quantizer = default_8bit_quantizers.Default8BitConvTransposeWeightsQuantizer(
+    )
+
+
 def _types_dict():
   return {
       'Default8BitQuantizeConfig':

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

@@ -41,3 +41,20 @@ def build(self, tensor_shape, name, layer):
         trainable=False)
 
     return {'min_var': min_weight, 'max_var': max_weight}
+
+
+class Default8BitConvTransposeWeightsQuantizer(quantizers.LastValueQuantizer):
+  """Quantizer for handling weights in Conv2DTranspose layers."""
+
+  def __init__(self):
+    """Construct LastValueQuantizer with params specific for TFLite Conv2DTranpose."""
+
+    super(Default8BitConvTransposeWeightsQuantizer, self).__init__(
+        num_bits=8, per_axis=False, symmetric=True, narrow_range=True)
+
+  def __call__(self, inputs, training, weights, **kwargs):
+    outputs = tf.transpose(inputs, (0, 1, 3, 2))
+    outputs = super(Default8BitConvTransposeWeightsQuantizer,
+                    self).__call__(outputs, training, weights, **kwargs)
+    outputs = tf.transpose(outputs, (0, 1, 3, 2))
+    return outputs

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

@@ -153,6 +153,13 @@ def _get_upsampling2d_bilinear_model(self):
     x = tf.keras.layers.UpSampling2D(size=(1, 5), interpolation='bilinear')(i)
     return tf.keras.Model(i, x)
 
+  def _get_conv2d_transpose_model(self):
+    i = tf.keras.Input(shape=(32, 32, 3))
+    x = tf.keras.layers.Conv2DTranspose(
+        2, kernel_size=(3, 3), strides=(2, 2))(
+            i)
+    return tf.keras.Model(i, x)
+
   @parameterized.parameters([
       _get_single_conv_model, _get_single_dense_model,
       _get_single_conv_relu_model, _get_stacked_convs_model,
@@ -165,6 +172,7 @@ def _get_upsampling2d_bilinear_model(self):
       # TODO(tfmot): There are gaps between ResizeBilinear with FakeQuant and
       # TFLite quantized ResizeBilinear op. It has a bit more quantization
       # error than other ops in this test now.
+      _get_conv2d_transpose_model,
   ])
   def testModelEndToEnd(self, model_fn):
     # 1. Check whether quantized model graph can be constructed.

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

@@ -298,7 +298,6 @@ class QuantizeFullIntegerModelTest(tf.test.TestCase, parameterized.TestCase):
           layers.UpSampling3D,
           # Not done since not registered since not per-axis yet.
           layers.Conv1D,
-          layers.Conv2DTranspose,
       ]
   ])
   def testQuantizeSingleLayer_ProducesFullIntegerModel_TF2(