将tensorflow的ckpt模型转化为pb模型

标签:tensorflow
作者:炼己者


本博客所有内容以学习、研究和分享为主,如需转载,请联系本人,标明作者和出处,并且是非商业用途,谢谢!


1.摘要

为什么好好的ckpt模型我要把它转为pb模型呢?因为我打算把python版本的代码转化为c++版本,最好的方法就是把python训练好的模型直接移植过去。但是ckpt模型不能用,所以我要想办法把训练好的ckpt模型转化为pb模型,然后再用c++调用这个模型。

2.ckpt转化为pb

我们用tensorflow训练模型,一般是用tf.train.Saver()保存模型,然后得到多个文件,一般长这个样子

将tensorflow的ckpt模型转化为pb模型_第1张图片

这四个文件主要是记录了神经网络的网络结构以及这个结构中涉及到的权重参数等内容。

具体怎么转呢?详细代码看ckpt转化为pb
点进去之后,直接clone即可。然后可以得到以下文件

checkpoints里面装着模型文件,点进去会发现里面有一个textcnn文件夹,这里面装着ckpt模型的四个文件。通过convert_ckpt_to_pb.py可以得到pb模型文件,这个文件也保存在checkpoints文件夹中,是frozen_model.pb。

def freeze_graph(input_checkpoint,output_graph):
    '''
    :param input_checkpoint:
    :param output_graph: PB模型保存路径
    :return:
    '''
    # 指定输出的节点名称,该节点名称必须是原模型中存在的节点
    # 直接用最后输出的节点,可以在tensorboard中查找到,tensorboard只能在linux中使用
    output_node_names = "score/output"
    saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
    graph = tf.get_default_graph() # 获得默认的图
    input_graph_def = graph.as_graph_def()  # 返回一个序列化的图代表当前的图
 
    with tf.Session() as sess:
        saver.restore(sess, input_checkpoint) #恢复图并得到数据
        output_graph_def = graph_util.convert_variables_to_constants(  # 模型持久化,将变量值固定
            sess=sess,
            input_graph_def=input_graph_def,# 等于:sess.graph_def
            output_node_names=output_node_names.split(","))# 如果有多个输出节点,以逗号隔开
 
        with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
            f.write(output_graph_def.SerializeToString()) #序列化输出
        print("%d ops in the final graph." % len(output_graph_def.node)) #得到当前图有几个操作节点

代码很简单,里面你唯一要改动的就是output_node_names,指定的最后一层输出节点名称,这个是你自己设定的。怎么找到它呢?首先你要去查看你的代码,
下面的代码是我自己定义的,你可以点击 训练ckpt模型,clone到本地,然后你会找到cnn_model.py这个文件,这里面就是定义着我的cnn模型

        with tf.name_scope("score"):
            # 全连接层,后面接dropout以及relu激活
            fc = tf.layers.dense(gmp, self.config.hidden_dim, name='fc1')
            fc = tf.contrib.layers.dropout(fc, self.keep_prob)
            fc = tf.nn.relu(fc)

            # 分类器
            self.logits = tf.layers.dense(fc, self.config.num_classes, name='fc2')
            self.y_pred_cls = tf.argmax(tf.nn.softmax(self.logits), 1,name='output')  # 预测类别

我一开始也没明白,然后翻阅大量资料才发现的,这个要自己改动,记得是最后输出的那层,你要自己把这个节点命名出来。另外怎么才能确定你得图结构里有这个节点呢?去查看tensorboard。不知道为啥在windows下得到的tensorboard查看不了,得到的是乱码文件。在linux下就不会,所以我把代码放到linux环境下跑,就可以得到tensoboard了。现在要跑的这个代码是指你自己训练ckpt模型的代码,如果你的tensorboard可以查看,那么就不用重新到linux环境下跑了。

查看tensorboard

tensorboard --logdir = “保存tensorboard的绝对路径”

敲入上面的命令,然后就可以得到一个网址,把这个网址复制到浏览器上打开,就可以得到图的结构,然后你点开看看,有没有output这个节点,也可以顺便看一下你自己的网络图
查看tensorboard的方法可以看这篇博客,TensorBoard:计算图的查看

有的话,把节点写入output_node_names,改这一个就行了。
改好之后,运行这个文件,便可以得到pb模型了。

3.pb模型文件测试数据

然后怎么用pb模型像ckpt模型那样来测试呢?c++版本的我还没开始做,但是Python版本的可以跑通测试了。这意味着我离c++跑通更进一步了。
这个具体去看pb_test.py文件,里面有很完整的注释。网上很多都是针对图像的资料,我这里是针对的文本的,相信对大家会有所帮助!
另外注意tensorflow是一个batch一个batch的逐批次输入,一次性输入的话会报内存错误!

4.总结

这两份代码大家可以对着看,很相似,理解怎么训练出来的ckpt模型,又是怎么把ckpt模型转化为pb模型的,怎么用pb模型去测试数据,代码很好理解的。有不懂的欢迎在博客下面评论提问,我们一起交流解决
等后期我用c++调用pb模型成功后,再来和大家分享!!!

以下是我所有文章的目录,大家如果感兴趣,也可以前往查看
戳右边:打开它,也许会看到很多对你有帮助的文章

你可能感兴趣的:(将tensorflow的ckpt模型转化为pb模型)