Profiling a Keras-TensorFlow model

keras-tensorflow-model-profiling.py

import tensorflow as tf
from tensorflow.python.client import timeline
from keras import backend as K

run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()

model = ...  # A Keras model

fn = K.function(model.inputs, model.outputs, options=run_options, run_metadata=run_metadata)

for i in range(4):
    x = K.variable(...)
    fn([x])
    tl = timeline.Timeline(run_metadata.step_stats)
    ctf = tl.generate_chrome_trace_format()
    with open('timeline_%d.json' % (i), 'w') as f:
        f.write(ctf)