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)

巴特西

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

最新文章

热门文章