-
Notifications
You must be signed in to change notification settings - Fork 24
/
export_model.py
77 lines (57 loc) · 2.52 KB
/
export_model.py
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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
"""This script will train and export a linear regression model into .pb, which
will is to be served by tensorflow serving.
"""
import os
import sys
import tensorflow as tf
import numpy as np
tf.app.flags.DEFINE_integer('training_iteration', 300,
'number of training iterations.')
tf.app.flags.DEFINE_integer('model_version', 1, 'version number of the model.')
tf.app.flags.DEFINE_string('work_dir', '', 'Working directory.')
FLAGS = tf.app.flags.FLAGS
def main(_):
sess = tf.InteractiveSession()
x = tf.placeholder('float', shape=[None, 3])
y_ = tf.placeholder('float', shape=[None, 1])
w = tf.get_variable('w', shape = [3,1], initializer = tf.truncated_normal_initializer)
b = tf.get_variable('b', shape = [1], initializer = tf.zeros_initializer)
sess.run(tf.global_variables_initializer())
y = tf.matmul(x, w) + b
ms_loss = tf.reduce_mean((y - y_)**2)
train_step = tf.train.GradientDescentOptimizer(0.005).minimize(ms_loss)
train_x = np.random.randn(1000, 3)
# let the model learn the equation of y = x1 * 1 + x2 * 2 + x3 * 3
train_y = np.sum(train_x * np.array([1,2,3]) + np.random.randn(1000, 3) / 100, axis = 1).reshape(-1, 1)
train_loss = []
for _ in range(FLAGS.training_iteration):
loss, _ = sess.run([ms_loss, train_step], feed_dict={x: train_x, y_: train_y})
train_loss.append(loss)
print('Training error %g' % loss)
print('Done training!')
# Export model
export_path_base = FLAGS.work_dir
export_path = os.path.join(
tf.compat.as_bytes(export_path_base),
tf.compat.as_bytes(str(FLAGS.model_version)))
print('Exporting trained model to', export_path)
builder = tf.saved_model.builder.SavedModelBuilder(export_path)
tensor_info_x = tf.saved_model.utils.build_tensor_info(x)
tensor_info_y = tf.saved_model.utils.build_tensor_info(y)
prediction_signature = (
tf.saved_model.signature_def_utils.build_signature_def(
inputs={'input': tensor_info_x},
outputs={'output': tensor_info_y},
method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME))
legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map={
'prediction':
prediction_signature,
},
legacy_init_op=legacy_init_op)
builder.save()
print('Done exporting!')
if __name__ == '__main__':
tf.app.run()