从 google 云机器学习引擎本地加载保存的 tensorflow 模型 .pb

Locally load saved tensorflow model .pb from google cloud machine learning engine

我想使用我在网上训练的 tensorflow 模型,然后 运行 使用我分发的 python 程序在本地使用它。

训练后,我得到一个目录 /model,其中包含两个文件 /saved_model.pb 和一个文件夹 /variables。在本地部署它的最简单方法是什么?

我正在关注 here 以部署冻结模型,但我无法完全阅读 .pb。我直接下载 saved_model.pb 到我的工作并尝试

with tf.gfile.GFile("saved_model.pb", "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

google.protobuf.message.DecodeError: Truncated message.

查看 SO ,他们建议了一条不同的路线。

with tf.gfile.GFile("saved_model.pb", "rb") as f:
    proto_b=f.read()
    graph_def = tf.GraphDef()
    text_format.Merge(proto_b, graph_def) 

builtins.TypeError: a bytes-like object is required, not 'str'

自从

我觉得这很混乱
type(proto_b)
<class 'bytes'>
type(graph_def)
<class 'tensorflow.core.framework.graph_pb2.GraphDef'>

为什么会报错,字符串也不是?

部署云训练模型的最佳方式是什么?

完整代码

import tensorflow as tf
import sys
from google.protobuf import text_format


# change this as you see fit
#image_path = sys.argv[1]
image_path="test.jpg"

# Read in the image_data
image_data = tf.gfile.FastGFile(image_path, 'rb').read()

# Loads label file, strips off carriage return
label_lines = [line.rstrip() for line 
               in tf.gfile.GFile("dict.txt")]

# Unpersists graph from file
with tf.gfile.GFile("saved_model.pb", "rb") as f:
    proto_b=f.read()
    graph_def = tf.GraphDef()
    text_format.Merge(proto_b, graph_def) 

with tf.Session() as sess:
    # Feed the image_data as input to the graph and get first prediction
    softmax_tensor = sess.graph.get_tensor_by_name('conv1/weights:0')

    predictions = sess.run(softmax_tensor, \
                           {'DecodeJpeg/contents:0': image_data})

    # Sort to show labels of first prediction in order of confidence
    top_k = predictions[0].argsort()[-len(predictions[0]):][::-1]

    for node_id in top_k:
        human_string = label_lines[node_id]
        score = predictions[0][node_id]
        print('%s (score = %.5f)' % (human_string, score))

您部署到 CloudML Engine 服务的模型格式是 SavedModel. Loading a SavedModel in Python is fairly simple using the loader 模块:

import tensorflow as tf

with tf.Session(graph=tf.Graph()) as sess:
   tf.saved_model.loader.load(
       sess,
       [tf.saved_model.tag_constants.SERVING],
       path_to_model)

进行推理,你的代码几乎是正确的;您将需要确保您正在向 session.run 提供一批,因此只需将 image_data 包装在列表中:

# Feed the image_data as input to the graph and get first prediction
softmax_tensor = sess.graph.get_tensor_by_name('conv1/weights:0')

predictions = sess.run(softmax_tensor, \
                       {'DecodeJpeg/contents:0': [image_data]})

# Sort to show labels of first prediction in order of confidence
top_k = predictions[0].argsort()[-len(predictions[0]):][::-1]

for node_id in top_k:
    human_string = label_lines[node_id]
    score = predictions[0][node_id]
    print('%s (score = %.5f)' % (human_string, score))

(请注意,根据您的图表,将您的 input_data 包装在列表中可能会增加您的预测张量的排名,您需要相应地调整代码)。