tf导出pb文件,以及如何使用pb文件

先罗列出来代码,有时间再解释

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



def export_model(input_checkpoint, output_graph):
    #这个可以加载saver的模型
    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:
        sess.run(tf.global_variables_initializer())
        
        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=['softmax_linear/softmax_linear','Cast_1'])# 如果有多个输出节点,以逗号隔开这个是重点,输入和输出的参数都需要在这里记录

        with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
            f.write(output_graph_def.SerializeToString()) #序列化输出
        

export_model('E:\\python\\image\\code2\\model_10\\model.ckpt',"E:\\python\\image\\code2\\model_10\\model.pb")

使用的代码

import os
import numpy as np
import tensorflow as tf
import model_new
from PIL import Image
import matplotlib.pyplot as plt
import csv
import shutil
from tensorflow.python.platform import gfile

def get_one_image(img_dir):
        
        image = Image.open(img_dir)
        
        image = image.resize((128,128))
        image = np.array(image)

        return image, img_dir


def test_model(model_path, img_path):
    image_array,img_dir = get_one_image( img_path)
    image = tf.cast(image_array,tf.float32)
    #image = tf.image.per_image_standardization(image)
    image = tf.reshape(image,[1,128,128,3])

    with tf.Session() as sess:
       
        with gfile.FastGFile(model_path,'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())
            sess.graph.as_default()
            tf.import_graph_def(graph_def,name='')
        sess.run(tf.global_variables_initializer())
        
        input_x = sess.graph.get_tensor_by_name('Cast_1:0')
        out = sess.graph.get_tensor_by_name('softmax_linear/softmax_linear:0')
        ret = sess.run(out,  feed_dict={input_x: image.eval()})
        print(ret)
        



out_pb_path="E:\\python\\image\\code2\\model_10\\frozen_model.pb"
img_path = "E:\\python\\image\\code\\images\\0\\mmexport1540880139708.jpg"
test_model(out_pb_path,img_path)
posted @ 2018-11-13 14:20  bbird2018  阅读(4034)  评论(0编辑  收藏  举报