import tensorflow as tf
import os

# 采用1号GPU，单GPU可注释
os.environ["CUDA_VISIBLE_DEVICES"] = "1"

# 输入pb模型的路径
input_pb_path = "model/pb/pb_model.pb"
# 输出tflie模型的路径
out_tflie_path = "model/tflite/tflite_model.tflite"

# 模型输入节点
input_tensor_name = ["x"]
input_tensor_shape = {input_tensor_name[0]:[1,4]}
# 模型输出节点
out_tensor_name = ['o']

converter = tf.lite.TFLiteConverter.from_frozen_graph(
    input_pb_path,
    input_tensor_name,
    out_tensor_name,
    input_tensor_shape)
converter.allow_custom_ops=True
converter.post_training_quantize = True
tflite_model = converter.convert()
open(out_tflie_path, "wb").write(tflite_model)
print('转换结束')