38 lines · python
1"""Convert a saved model to tflite model.2 3Usage: python3 saved-model-to-tflite.py <mlgo saved_model_dir> <tflite dest_dir>4 5The <tflite dest_dir> will contain:6 model.tflite: this is the converted saved model7 output_spec.json: the output spec, copied from the saved_model dir.8"""9 10import tensorflow as tf11import os12import sys13from tf_agents.policies import greedy_policy14 15 16def main(argv):17 assert len(argv) == 318 sm_dir = argv[1]19 tfl_dir = argv[2]20 tf.io.gfile.makedirs(tfl_dir)21 tfl_path = os.path.join(tfl_dir, "model.tflite")22 converter = tf.lite.TFLiteConverter.from_saved_model(sm_dir)23 converter.target_spec.supported_ops = [24 tf.lite.OpsSet.TFLITE_BUILTINS,25 ]26 tfl_model = converter.convert()27 with tf.io.gfile.GFile(tfl_path, "wb") as f:28 f.write(tfl_model)29 30 json_file = "output_spec.json"31 src_json = os.path.join(sm_dir, json_file)32 if tf.io.gfile.exists(src_json):33 tf.io.gfile.copy(src_json, os.path.join(tfl_dir, json_file))34 35 36if __name__ == "__main__":37 main(sys.argv)38