将图形proto(pb/pbtxt)转换为SavedModel,以便在TensorFlow服务或Cloud ML Engine中使用

Question

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

由于我可以使用TensorFlow Android推理库推断该图表,我认为我可以使用Cloud ML Engine做同样的事情,但它似乎只适用于SavedModel模型.

如何在单个pb文件中简单地转换冻结/量化图以在ML引擎上使用？

Answer 1

事实证明,SavedModel为保存的图形提供了一些额外的信息.假设冻结图不需要资产,那么它只需要指定服务签名.

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

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()