convert_to_tflite.py

import sys
import os

import tensorflow as tf


def usage():
    proggie = os.path.basename(sys.argv[0])
    print(f"""\
Usage: {proggie} <input_saved_model_dir> <output>

Example: {proggie} output_inference_graph_v1.pb/saved_model converted_model.tflite""")


if __name__ == "__main__":
    print(f"TF Version: ${tf.__version__}")
    if len(sys.argv) != 3 or '-h' in sys.argv or '--help' in sys.argv:
        usage()
        sys.exit(1)

    saved_model_dir, outfile = sys.argv[1:]
    converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir,
                                                         signature_keys=['serving_default'])
    converter.optimizations = [tf.lite.Optimize.DEFAULT]
    converter.experimental_new_converter = True
    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS,
                                           tf.lite.OpsSet.SELECT_TF_OPS]
    tflite_model = converter.convert()
    with tf.io.gfile.GFile(outfile, 'wb') as f:
        f.write(tflite_model)
    print(f"{outfile} written")