-
Notifications
You must be signed in to change notification settings - Fork 3
Expand file tree
/
Copy pathexport_model.py
More file actions
32 lines (25 loc) · 965 Bytes
/
export_model.py
File metadata and controls
32 lines (25 loc) · 965 Bytes
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
import os
import tensorflow as tf
import tflearn
from tensorflow.contrib.session_bundle import exporter
from train import MODEL_NAME, MODEL_PATH, resnet
pjoin = os.path.join
def export_model(path):
net = resnet()
g = net.graph
X = g.get_operation_by_name('InputData/X')
saver = tf.train.Saver()
with tf.Session() as sess:
model = tflearn.DNN(net, session=sess)
# model.load(pjoin(MODEL_PATH, MODEL_NAME))
saver.restore(sess, pjoin(MODEL_PATH, MODEL_NAME))
model_exporter = exporter.Exporter(saver)
model_exporter.init(
sess.graph.as_graph_def(),
named_graph_signatures={
'inputs': exporter.generic_signature({'x': X}),
'outputs': exporter.generic_signature({'y': net})})
model_exporter.export(path, tf.constant(200), sess)
print 'Successfully exported model to %s' % path
if __name__ == '__main__':
export_model('export')