Converting TensorFlow 2 Detection Model Zoo Models to TFLite

1. Introduction

There is a repository called TensorFlow 2 Detection Model Zoo. Various pre-trained object detection models based on TensorFlow v2 are committed here, but since there is no mention of a formal or established procedure anywhere, I converted them appropriately on my own. I will share the steps. This is a completely throwaway article because it will be over in a flash.
https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md

2. Procedure

Download and extract the pre-trained model.

$ wget http://download.tensorflow.org/models/object_detection/tf2/20200711/centernet_resnet50_v2_512x512_kpts_coco17_tpu-8.tar.gz
$ tar -zxvf centernet_resnet50_v2_512x512_kpts_coco17_tpu-8.tar.gz
$ rm centernet_resnet50_v2_512x512_kpts_coco17_tpu-8.tar.gz
$ cd centernet_resnet50_v2_512x512_kpts_coco17_tpu-8

Check the input and output names and shapes of the saved_model included when you downloaded and extracted it. Run the following command:

$ saved_model_cli show \
  --dir saved_model \
  --tag_set serve \
  --signature_def serving_default

The given SavedModel SignatureDef contains the following input(s):
  inputs['input_tensor'] tensor_info:
      dtype: DT_UINT8
      shape: (1, -1, -1, 3)
      name: serving_default_input_tensor:0
The given SavedModel SignatureDef contains the following output(s):
  outputs['detection_boxes'] tensor_info:
      dtype: DT_FLOAT
      shape: (1, 100, 4)
      name: StatefulPartitionedCall:0
  outputs['detection_classes'] tensor_info:
      dtype: DT_FLOAT
      shape: (1, 100)
      name: StatefulPartitionedCall:1
  outputs['detection_keypoint_scores'] tensor_info:
      dtype: DT_FLOAT
      shape: (1, 100, 17)
      name: StatefulPartitionedCall:2
  outputs['detection_keypoints'] tensor_info:
      dtype: DT_FLOAT
      shape: (1, 100, 17, 2)
      name: StatefulPartitionedCall:3
  outputs['detection_scores'] tensor_info:
      dtype: DT_FLOAT
      shape: (1, 100)
      name: StatefulPartitionedCall:4
  outputs['num_detections'] tensor_info:
      dtype: DT_FLOAT
      shape: (1)
      name: StatefulPartitionedCall:5
Method name is: tensorflow/serving/predict

Create the following logic and save it as test.py. Use something called a concrete function to fix the input resolution of the saved_model from [1, -1, -1, 3] to [1, 320, 320, 3]. Feel free to change the target resolution as you like. TFLite conversion will fail if dimensions other than the batch size remain undefined (-1). Note that if there are multiple input Tensors, simply list them in the input_shapes = [[1,320,320,3]] part like input_shapes = [[1,320,320,3],[1,10,10,1]].

### tf-nightly==2.6.0-dev20210430

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2_as_graph

# Load saved_model and change input shape
# https://github.com/tensorflow/tensorflow/issues/30180#issuecomment-505959220
model = tf.saved_model.load('saved_model')
input_shapes = [[1,320,320,3]]

concrete_func = \
    model.signatures[tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
concrete_func_input_tensors = \
    [tensor for tensor in concrete_func.inputs if tensor.dtype != tf.resource and not 'unknown' in tensor.name]

for conc_input, def_input in zip(concrete_func_input_tensors, input_shapes):
    print('Before changing the input shape', conc_input)
    conc_input.set_shape(def_input)
    print('After  changing the input shape', conc_input)

converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
converter.target_ops = \
    [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
tflite_model = converter.convert()
open("test.tflite", "wb").write(tflite_model)

Clean up the TensorFlow environment. If you don't want to break your local environment, please use this Docker container. You can quickly use an environment that you can pull from Docker Hub and destroy as many times as you like.
https://github.com/PINTO0309/tflite2tensorflow

$ sudo pip3 uninstall -y tensorboard-plugin-wit tb-nightly tensorboard \
                      tf-estimator-nightly tensorflow-gpu \
                      tensorflow tf-nightly tensorflow_estimator

Install tf-nightly.

$ sudo pip3 install tf-nightly==2.6.0-dev20210430

Run the conversion script.

$ python3 test.py

Done. It's easy, right? You can see that a file named test.tflite has been generated.