将图形原型 (pb/pbtxt) 转换为 SavedModel 以在 TensorFlow Serving 或 Cloud ML Engine 中使用

问题描述

我一直在关注 TensorFlow for Poets 2我训练过的模型上的 codelab，并创建了一个带有嵌入权重的冻结的量化图.它被捕获在一个文件中 - 比如 my_quant_graph.pb.

I've been following the TensorFlow for Poets 2 codelab on a model I've trained, and have created a frozen, quantized graph with embedded weights. It's captured in a single file - say my_quant_graph.pb.

因为我可以通过 TensorFlow Android 推理库使用该图进行推理 很好，我以为我可以用 Cloud ML Engine 做同样的事情，但它似乎只适用于 SavedModel 模型.

Since I can use that graph for inference with the TensorFlow Android inference library just fine, I thought I could do the same with Cloud ML Engine, but it seems it only works on a SavedModel model.

如何简单地将单个 pb 文件中的冻结/量化图形转换为在 ML 引擎上使用?

How can I simply convert a frozen/quantized graph in a single pb file to use on ML engine?

推荐答案

结果是 SavedModel 提供了一些关于保存的图形的额外信息.假设冻结图不需要资产，那么它只需要指定一个服务签名.

It turns out that a SavedModel provides some extra info around a saved graph. Assuming a frozen graph doesn't need assets, then it needs only a serving signature specified.

这是我运行的 Python 代码，用于将图形转换为 Cloud ML 引擎接受的格式.注意我只有一对输入/输出张量.

Here's the python code I ran to convert my graph to a format that Cloud ML engine accepted. Note I only have a single pair of input/output tensors.

import tensorflow as tf
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants

export_dir = './saved'
graph_pb = 'my_quant_graph.pb'

builder = tf.saved_model.builder.SavedModelBuilder(export_dir)

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

sigs = {}

with tf.Session(graph=tf.Graph()) as sess:
    # name="" is important to ensure we don't get spurious prefixing
    tf.import_graph_def(graph_def, name="")
    g = tf.get_default_graph()
    inp = g.get_tensor_by_name("real_A_and_B_images:0")
    out = g.get_tensor_by_name("generator/Tanh:0")

    sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] =
        tf.saved_model.signature_def_utils.predict_signature_def(
            {"in": inp}, {"out": out})

    builder.add_meta_graph_and_variables(sess,
                                         [tag_constants.SERVING],
                                         signature_def_map=sigs)

builder.save()

这篇关于将图形原型 (pb/pbtxt) 转换为 SavedModel 以在 TensorFlow Serving 或 Cloud ML Engine 中使用的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！