TensorFlow/test/createconstantgraph.py

0001 # coding: utf-8
0002
0003 """
0004 Test script to create a simple graph for testing purposes at bin/data and save it with all
0005 variables converted to constants to reduce its memory footprint.
0006
0007 https://www.tensorflow.org/api_docs/python/tf/graph_util/convert_variables_to_constants
0008 """
0009
0010 import os
0011 import sys
0012
0013 import cmsml
0014
0015
0016 # get tensorflow and work with the v1 compatibility layer
0017 tf, tf1, tf_version = cmsml.tensorflow.import_tf()
0018 tf = tf1
0019 tf.disable_eager_execution()
0020
0021 # prepare the datadir
0022 if len(sys.argv) >= 2:
0023     datadir = sys.argv[1]
0024 else:
0025     thisdir = os.path.dirname(os.path.abspath(__file__))
0026     datadir = os.path.join(os.path.dirname(thisdir), "bin", "data")
0027
0028 # create the graph
0029 x_ = tf.placeholder(tf.float32, [None, 10], name="input")
0030 scale_ = tf.placeholder(tf.float32, name="scale")
0031
0032 W = tf.Variable(tf.ones([10, 1]))
0033 b = tf.Variable(tf.ones([1]))
0034 h = tf.add(tf.matmul(x_, W), b)
0035 y = tf.multiply(h, scale_, name="output")
0036
0037
0038 # Setup the script to run on CPU only
0039 config = tf.ConfigProto(
0040         device_count = {'GPU': 0}
0041     )
0042 sess = tf.Session(config=config)
0043 sess.run(tf.global_variables_initializer())
0044
0045 print(sess.run(y, feed_dict={scale_: 1.0, x_: [range(10)]})[0][0])
0046
0047 # write it
0048 graph_path = os.path.join(datadir, "constantgraph.pb")
0049 outputs = ["output"]
0050 cmsml.tensorflow.save_frozen_graph(graph_path, sess, output_names=outputs, variables_to_constants=True)