1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import tensorflow as tf
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
from google.protobuf import text_format
import os
export_dir = 'inference/pb2saved'
graph_pb = '/Users/edz/PycharmProjects/serviceRec/train/saved_model_test/1640591747/saved_model.pb'
if os.path.exists(export_dir):
os.rmdir(export_dir)
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()
print(sess.graph.get_name_scope())
print(sess.graph.get_all_collection_keys())
print(sess.graph.get_operations())
# input_ids = sess.graph.get_tensor_by_name(
# "input_ids:0")
# input_mask = sess.graph.get_tensor_by_name(
# "input_mask:0")
# segment_ids = sess.graph.get_tensor_by_name(
# "segment_ids:0")
# probabilities = g.get_tensor_by_name("loss/pred_prob:0")
# sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \
# tf.saved_model.signature_def_utils.predict_signature_def(
# {
# "input_ids": input_ids,
# "input_mask": input_mask,
# "segment_ids": segment_ids
# }, {
# "probabilities": probabilities
# })
# builder.add_meta_graph_and_variables(sess,
# [tag_constants.SERVING],
# signature_def_map=sigs)
# builder.save()