bilinear upsampling layer cannot be 8-bit quantized correctly

Describe the bug

Hello, I want to do full 8-bit quantization(input, weight all 8-bit) to the network with a bilinear upsampling layer. The fake QAT result in validation set is closed to FP32 result, but when I converted the QAT model to full 8-bit tflite model, the result in validation set decreased significantly, almost no accuracy. So I wonder whether I made a mistake or tflite doesn't support int8 bilinear upsampling correctly.

System information

TensorFlow version (installed from source or binary): 2.5.0

TensorFlow Model Optimization version (installed from source or binary): 0.7.2

Python version: 3.7.13

Describe the current behavior

I aim to do 8-bit quantization to the network below, which has a bilinear upsample branch

def upsambiskip(scale=3, in_channels=3, num_fea=28, m=4, out_channels=3):
    inp = Input(shape=(None, None, 3)) 
    upsampled_inp=UpSampling2D(size=(3,3),data_format=None,interpolation='bilinear')(inp)

    x = Conv2D(num_fea, 3, padding='same', activation='relu', kernel_initializer=glorot_normal(), bias_initializer='zeros')(inp)

    for i in range(m):
        x = Conv2D(num_fea, 3, padding='same', activation='relu', kernel_initializer=glorot_normal(), bias_initializer='zeros')(x)

    x = Conv2D(out_channels*(scale**2), 3, padding='same', kernel_initializer=glorot_normal(), bias_initializer='zeros')(x)
    
    depth_to_space = Lambda(lambda x: tf.nn.depth_to_space(x, scale))
    out = depth_to_space(x)
    x = Add()([upsampled_inp, out])
    x = Conv2D(3, 3, padding='same', kernel_initializer=glorot_normal(), bias_initializer='zeros')(x)

    clip_func = Lambda(lambda x: K.clip(x, 0., 255.))
    out = clip_func(x)
    
    return Model(inputs=inp, outputs=out)

I had done FP32 training before, then I loaded FP32 checkpoint and did Quantization-aware training like below

class NoOpQuantizeConfig(tfmot.quantization.keras.QuantizeConfig):
    def get_weights_and_quantizers(self, layer):
        return []
    def get_activations_and_quantizers(self, layer):
        return []
    def set_quantize_weights(self, layer, quantize_weights):
        pass
    def set_quantize_activations(self, layer, quantize_anctivations):
        pass
    def get_output_quantizers(self, layer):
        return []
    def get_config(self):
        return {}

def ps_quantization(self, layer):
    # lambda not quantization
    if 'lambda' in layer.name :
        return tfmot.quantization.keras.quantize_annotate_layer(layer, quantize_config=NoOpQuantizeConfig())
    return layer


# create fp32 model and load weight
p_model = create_model(args['networks'])
    
lg.info('Start copying weights and annotate Lambda layer...')
annotate_model = tf.keras.models.clone_model(
            p_model,
            clone_function=self.ps_quantization
            )
lg.info('Start annotating other parts of model...')
annotate_model = tfmot.quantization.keras.quantize_annotate_model(annotate_model)
lg.info('Creating quantize-aware model...')
depth_to_space = Lambda(lambda x: tf.nn.depth_to_space(x, 3))
with tfmot.quantization.keras.quantize_scope({'NoOpQuantizeConfig': NoOpQuantizeConfig, 'depth_to_space': depth_to_space, 'tf': tf}):
    self.model = tfmot.quantization.keras.quantize_apply(annotate_model)
    
# training...

the QAT model was correct, and the performance in validation set was closed to FP32's. Then I converted QAT model to tflite

def qat_quantize(model_path,output_path):

    fakeqmodel=tf.keras.models.load_model(model_path)
    print('fake quantize model validate..')

    # validate fake QAT model, and the performance is OK
    load_validate(fakeqmodel,save_path=None)

    # convert QAT model and store int8 tflite model
    converter = tf.lite.TFLiteConverter.from_keras_model(fakeqmodel)
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    ### converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
    converter.experimental_new_converter=True
    converter.experimental_new_quantizer=True
    converter.inference_input_type = tf.uint8
    converter.inference_output_type = tf.uint8
    quant_tf_model = converter.convert()
    with open(output_path, 'wb') as f:
        f.write(quant_tf_model)

I loaded tflite model and evaluated as below, but the performance decrease significantly.

def evaluate(quantized_model_path, save_path):

    interpreter = tf.lite.Interpreter(model_path=quantized_model_path,num_threads=32)
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    IS, IZ = input_details[0]['quantization']
    OS, OZ = output_details[0]['quantization']
    print('Input Scale: {}, Zero Point: {}'.format(IS, IZ))
    print('Output Scale: {}, Zero Point: {}'.format(OS, OZ))
    psnr = 0.0
    for i in range(801, 901):
        lr_path = 'data/DIV2K/DIV2K_train_LR_bicubic/X3_pt/0{}x3.pt'.format(i)
        with open(lr_path, 'rb') as f:
            lr = pickle.load(f)
        h, w, c = lr.shape
        lr = np.expand_dims(lr, 0).astype(np.float32)
        # ##lr = np.round(lr/IS+IZ).astype(np.uint8)
        lr = lr.astype(np.uint8)

        hr_path = 'data/DIV2K/DIV2K_train_HR_pt/0{}.pt'.format(i)
        with open(hr_path, 'rb') as f:
            hr = pickle.load(f)
        hr = np.expand_dims(hr, 0).astype(np.float32)
        interpreter.resize_tensor_input(input_details[0]['index'], lr.shape)
        interpreter.allocate_tensors()
        interpreter.set_tensor(input_details[0]['index'], lr)
        interpreter.invoke()

        sr = interpreter.get_tensor(output_details[0]['index'])
        #sr = np.clip(np.round((sr.astype(np.float32)-OZ)*OS), 0, 255)
        sr = np.clip(sr, 0, 255)
        b, h, w, c = sr.shape
        # save image
        if save_path is not None:
            save_name = osp.join(save_path, '{:04d}x3.png'.format(i))
            cv2.imwrite(save_name, cv2.cvtColor(sr.squeeze().astype(np.uint8), cv2.COLOR_RGB2BGR))

        mse = np.mean((sr[:, 1:h-1, 1:w-1, :].astype(np.float32) - hr[:, 1:h-1, 1:w-1, :].astype(np.float32)) ** 2)
        singlepsnr =  20. * math.log10(255. / math.sqrt(mse))
        print('[{}]/[100]: {}'.format(i, singlepsnr))
        psnr += singlepsnr
    print(psnr / 100)

Describe the expected behavior
I thought the wrong behavior was due to blinear upsampling layer, because when I removed upsample layer(just plain net without skip connection), or replaced upsample skip connection with a conv3x3 skip, I could get the right int8 tflite performance closed to the QAT performance. So I want to know how to quantize bilinear upsampling layer correctly?

other info
The tflite model is below

Start with the UpSampling2D bilinear branch and the custom quantization configuration, then inspect the TFLiteConverter settings and the evaluation input/output handling. Compare the fake-QAT model with the converted model using the reported TensorFlow 2.5.0 and TensorFlow Model Optimization 0.7.2 setup. Done means identifying whether bilinear upsampling conversion or evaluation causes the accuracy loss and documenting or fixing the supported quantization path.