Add Upsampling2D supports.

There are two interpolation mode for UpSampling2D op. (nearest and bilinear)
ResizeNearest just passed numerical test, but ResizeBilinear has larger quantization error for now.

PiperOrigin-RevId: 340558630

Author: Xhark

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

@@ -97,7 +97,21 @@ class QuantizeRegistry(quantize_registry.QuantizeRegistry, _RNNHelper):
       _no_quantize(layers.Cropping2D),
       _no_quantize(layers.Cropping3D),
       # _no_quantize(layers.UpSampling1D),
-      # _no_quantize(layers.UpSampling2D),
+
+      # TODO(tfmot): Reduce the quantization errors for bilinear interpolation
+      # type for UpSampling2D op. UpSampling2D supports two interpolation types,
+      # nearest and bilinear. we convert the op to ResizeBilnear integer op on
+      # TFLite. This ResizeBilinear TFLite op only for input and output has the
+      # same quantization parameters. (scale and zero_point) To do that, The
+      # TFLite converter inserts quantization cast op right after the input to
+      # match quantization params for the output. Current QAT doesn’t consider
+      # this behavior yet, so now we have larger quantization errors than we
+      # expected. We have to add support for it on QAT or change the TFLite
+      # kernel op to support different quantization params for input and output.
+      # (Note that the nearest case just copies the number so there’s no more
+      # errors even if the quantization order is different.)
+      _QuantizeInfo(layers.UpSampling2D, [], [], True),
+
       # _no_quantize(layers.UpSampling3D),
       _no_quantize(layers.ZeroPadding1D),
       _no_quantize(layers.ZeroPadding2D),

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

@@ -143,13 +143,28 @@ def _get_sepconv1d_stacked_model(self):
     x = tf.keras.layers.ReLU()(x)
     return tf.keras.Model(i, x)
 
+  def _get_upsampling2d_nearest_model(self):
+    i = tf.keras.Input(shape=(32, 32, 3))
+    x = tf.keras.layers.UpSampling2D(size=(3, 4), interpolation='nearest')(i)
+    return tf.keras.Model(i, x)
+
+  def _get_upsampling2d_bilinear_model(self):
+    i = tf.keras.Input(shape=(1, 3, 1))
+    x = tf.keras.layers.UpSampling2D(size=(1, 5), interpolation='bilinear')(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,
       _get_conv_bn_relu_model, _get_depthconv_bn_relu_model,
       _get_separable_conv2d_model,
       _get_sepconv1d_bn_model, _get_sepconv1d_bn_relu_model,
-      _get_sepconv1d_stacked_model
+      _get_sepconv1d_stacked_model,
+      _get_upsampling2d_nearest_model,
+      # _get_upsampling2d_bilinear_model
+      # 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.
   ])
   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

@@ -141,6 +141,11 @@ class QuantizeFullIntegerModelTest(tf.test.TestCase, parameterized.TestCase):
       }),
       (layers.UpSampling2D, {
           'input_shape': (4, 6, 1),
+          'interpolation': 'nearest',
+      }),
+      (layers.UpSampling2D, {
+          'input_shape': (4, 6, 1),
+          'interpolation': 'bilinear',
       }),
       (layers.UpSampling3D, {
           'input_shape': (4, 6, 1),
@@ -266,7 +271,6 @@ class QuantizeFullIntegerModelTest(tf.test.TestCase, parameterized.TestCase):
   @parameterized.parameters([
       l for l in _LAYER_PARAMS if l[0] not in [
           # Not done since TFLite converter doesn't support in TF2 yet.
-          layers.UpSampling2D,
           layers.Conv3D,
           layers.Conv3DTranspose,
           layers.AveragePooling3D,