import tensorflow as tf
import os
from tensorflow.python.framework import graph_util

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

# 输入ckpt模型路径
input_ckpt_path = 'model/ckpt/model-9000'
# 输出pb模型的路径
out_pb_path = "model/pb/pb_model.pb"

# 初始化
output_node_names = "o"
saver = tf.train.import_meta_graph(input_ckpt_path+'.meta', clear_devices=True)

with tf.Session() as sess:
    saver.restore(sess, input_ckpt_path)
    output_graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(
          sess=sess,
          input_graph_def=sess.graph_def,
          output_node_names=output_node_names.split(","))
    with tf.io.gfile.GFile(out_pb_path, "wb") as f:
        f.write(output_graph_def.SerializeToString())
print('转换结束')