tensorflow保存为pb格式模型

保存模型

if i % 200 == 0:
    print("After %d training step(s), loss on training batch is %g." % (step, loss_value))
    saver.save(sess, os.path.join(MODEL_SAVE_PATH, MODEL_NAME), global_step=global_step)
    #以上是常见的保存为ckpt格式模型,下面是保存为pb格式模型,加上下面这段代码即可
    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ["output"])#["output"]是关键,涉及到保存到哪一层operation
    #convert_variables_to_constants()方法,可以固化模型结构,将计算图中的变量取值以常量的形式保存,所以最好给神经网络最后一层操作命名以区分。
    with tf.gfile.FastGFile('model.pb', mode='wb') as f:
        f.write(constant_graph.SerializeToString())#将模型保存到save/model.pb文件,

你可能感兴趣的:(Python,TensorFlow,计算机视觉)