tf-profile.txt

tf-profile.txt

import tensorflow as tf


class Profiler(object):
    def __init__(self):
        super(Profiler, self).__init__()
        self.step = 0

    def profile_run(self, sess, fetches, feed_dict, filename='profile.txt'):
        profiler = tf.profiler.Profiler(sess.graph)
        run_metadata = tf.RunMetadata()
        sess.run(fetches, feed_dict=feed_dict,
                        options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE),
                        run_metadata=run_metadata)
        profiler.add_step(self.step, run_metadata)
        option_builder = tf.profiler.ProfileOptionBuilder
        opts = (option_builder(option_builder.time_and_memory()).
            with_step(self.step). # with -1, should compute the average of all registered steps.
            with_file_output(filename).
            select(['micros']).order_by('micros').
            build())
        profiler.profile_operations(options=opts)

        self.step += 1


profiler = Profiler()

learning_rate = 4e-3
cost = ...
optimizer = tf.train.RMSPropOptimizer(learning_rate).minimize(cost)
init = tf.global_variables_initializer()

with tf.Session() as sess:

    sess.run(init)

    for i, batch in enumerate(batches):

        profiler.profile_run(sess, [optimizer], feed_dict,
            filename='profile-step_{}.txt'.format(i))