Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2021-02-24 09:53:24

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 sess = tf.Session()
0038 sess.run(tf.global_variables_initializer())
0039 
0040 print(sess.run(y, feed_dict={scale_: 1.0, x_: [range(10)]})[0][0])
0041 
0042 # write it
0043 graph_path = os.path.join(datadir, "constantgraph.pb")
0044 outputs = ["output"]
0045 cmsml.tensorflow.save_graph(graph_path, sess, output_names=outputs, variables_to_constants=True)