Convert a graph proto (pb/pbtxt) to a SavedModel for use in TensorFlow Serving or Cloud ML Engine

前端 未结 1 1819
囚心锁ツ
囚心锁ツ 2020-11-28 08:40

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 singl

相关标签:
1条回答
  • 2020-11-28 09:46

    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.

    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()
    
    0 讨论(0)
提交回复
热议问题