tensorflow中ckpt转pb模型

参考代码：

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

def set_config():#设置GPU使用率
    # 控制使用率
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    # 假如有16GB的显存并使用其中的8GB:
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.3)
    config = tf.ConfigProto(gpu_options=gpu_options)
    # session = tf.Session(config=config)
    return config

#下面是你自定义的模型

inputs = tf.placeholder(tf.float32, shape=[None,width, height, 1], name='inputs')
labels = tf.placeholder(tf.float32, shape=[None,num_classes], name='labels')
pred=dense_net(inputs, num_classes,nb_blocks, growth_rate,weight_decay,compression)

model_path  = "./model/4_class.ckpt-1" #设置model的路径，因新版tensorflow会生成三个文件，只需写到数字前

cfg=set_config()

from tensorflow.python.saved_model import signature_constants, signature_def_utils, tag_constants, utils
save_path = 'model_pb/0'
with tf.Session(config=cfg) as sess:
    saver = tf.train.Saver()
    saver.restore(sess, model_path)
    model_signature = signature_def_utils.build_signature_def(
            inputs={"input": utils.build_tensor_info(inputs)},
            outputs={"pred": utils.build_tensor_info(pred)},
            method_name=signature_constants.PREDICT_METHOD_NAME)

    builder = tf.saved_model.builder.SavedModelBuilder(save_path)
    legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
    builder.add_meta_graph_and_variables(
              sess, [tag_constants.SERVING],
              clear_devices=True,
              signature_def_map={
                  signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                  model_signature,
              },
              legacy_init_op=legacy_init_op)
    builder.save()