python模型持久化_tensorflow的ckpt及pb模型持久化方式及转化详解

使用tensorflow训练模型的时候，模型持久化对我们来说非常重要。

如果我们的模型比较复杂，需要的数据比较多，那么在模型的训练时间会耗时很长。如果在训练过程中出现了模型不可预期的错误，导致训练意外终止，那么我们将会前功尽弃。为了解决这一问题，我们可以使用模型持久化(保存为ckpt文件格式)来保存我们在训练过程中的临时数据。、

如果我们训练出的模型需要提供给用户做离线预测，那么我们只需要完成前向传播过程。这个时候我们就可以使用模型持久化(保存为pb文件格式)来只保存前向传播过程中的变量并将变量固定下来，这时候用户只需要提供一个输入即可得到前向传播的预测结果。

ckpt和pb持久化方式的区别在于ckpt文件将模型结构与模型权重分离保存，便于训练过程;pb文件则是graph_def的序列化文件，便于发布和离线预测。官方提供freeze_grpah.py脚本来将ckpt文件转为pb文件。

CKPT模型持久化

首先定义前向传播过程;

声明并得到一个Saver;

使用Saver.save()保存模型;

# coding=UTF-8 支持中文编码格式

import tensorflow as tf

import shutil

import os.path

MODEL_DIR = "/home/zheng/PycharmProjects/ckptLoad/Models/"

MODEL_NAME = "model.ckpt"

#下面的过程你可以替换成CNN、RNN等你想做的训练过程，这里只是简单的一个计算公式

input_holder = tf.placeholder(tf.float32, shape=[1], name="input_holder") #输入占位符，并指定名字，后续模型读取可能会用的

W1 = tf.Variable(tf.constant(5.0, shape=[1]), name="W1")

B1 = tf.Variable(tf.constant(1.0, shape=[1]), name="B1")

_y = (input_holder * W1) + B1

predictions = tf.add(_y, 50, name="predictions") #输出节点名字，后续模型读取会用到，比50大返回true，否则返回false

init = tf.global_variables_initializer()

saver = tf.train.Saver() #声明saver用于保存模型

with tf.Session() as sess:

sess.run(init)

print "predictions : ", sess.run(predictions, feed_dict={input_holder: [10.0]}) #输入一个数据测试一下

saver.save(sess, os.path.join(MODEL_DIR, MODEL_NAME)) #模型保存

print("%d ops in the final graph." % len(tf.get_default_graph().as_graph_def().node)) #得到当前图有几个操作节点

predictions : [ 101.]

28 ops in the final graph.

注：代码含义请参考注释，需要注意的是可以自定义模型保存的路径

ckpt模型持久化使用起来非常简单，只需要我们声明一个tf.train.Saver&