def _merge_default_with_oplog(graph, op_log=None, run_meta=None):
"""Merge the tfprof default extra info with caller's op_log.
Args:
graph: tf.Graph.
op_log: OpLog proto.
run_meta: RunMetadata proto used to complete shape information.
Returns:
tmp_op_log: Merged OpLog proto.
"""
tmp_op_log = tfprof_log_pb2.OpLog()
logged_ops = _get_logged_ops(graph, run_meta)
if not op_log:
tmp_op_log.log_entries.extend(logged_ops.values())
else:
all_ops = dict()
for entry in op_log.log_entries:
all_ops[entry.name] = entry
for op_name, entry in six.iteritems(logged_ops):
if op_name in all_ops:
all_ops[op_name].types.extend(entry.types)
if entry.float_ops > 0 and all_ops[op_name].float_ops == 0:
all_ops[op_name].float_ops = entry.float_ops
else:
all_ops[op_name] = entry
tmp_op_log.log_entries.extend(all_ops.values())
return tmp_op_log
def _merge_default_with_oplog(graph,
op_log=None,
run_meta=None,
add_trace=True,
add_trainable_var=True):
"""Merge the tfprof default extra info with caller's op_log.
Args:
graph: tf.Graph.
op_log: OpLog proto.
run_meta: RunMetadata proto used to complete shape information.
add_trace: Whether to add op trace information.
add_trainable_var: Whether to assign tf.trainable_variables() op type
'_trainable_variables'.
Returns:
tmp_op_log: Merged OpLog proto.
"""
tmp_op_log = tfprof_log_pb2.OpLog()
logged_ops = _get_logged_ops(graph,
run_meta,
add_trace=add_trace,
add_trainable_var=add_trainable_var)
if not op_log:
tmp_op_log.log_entries.extend(logged_ops.values())
else:
all_ops = dict()
for entry in op_log.log_entries:
all_ops[entry.name] = entry
for op_name, entry in six.iteritems(logged_ops):
if op_name in all_ops:
all_ops[op_name].types.extend(entry.types)
if entry.float_ops > 0 and all_ops[op_name].float_ops == 0:
all_ops[op_name].float_ops = entry.float_ops
if entry.code_def.traces and not all_ops[
op_name].code_def.traces:
all_ops[op_name].code_def.MergeFrom(entry.code_def)
else:
all_ops[op_name] = entry
tmp_op_log.log_entries.extend(all_ops.values())
return tmp_op_log
def load_metadata(model_dir):
"""Loads RunMetadata, Graph and OpLog from files
"""
# Import RunMetadata
run_meta_path = os.path.join(model_dir, "metadata/run_meta")
run_meta = tf.RunMetadata()
if gfile.Exists(run_meta_path):
with gfile.GFile(run_meta_path, "rb") as file:
run_meta.MergeFromString(file.read())
print("Loaded RunMetadata from {}".format(run_meta_path))
else:
print(
"RunMetadata does not exist a {}. Skipping.".format(run_meta_path))
# Import Graph
graph_def_path = os.path.join(model_dir, "graph.pbtxt")
graph = tf.Graph()
if gfile.Exists(graph_def_path):
with graph.as_default():
_register_function_ops(CUSTOM_OP_FUNCTIONS)
graph_def = tf.GraphDef()
with gfile.GFile(graph_def_path, "rb") as file:
text_format.Parse(file.read(), graph_def)
tf.import_graph_def(graph_def, name="")
print("Loaded Graph from {}".format(graph_def_path))
else:
print("Graph does not exist a {}. Skipping.".format(graph_def_path))
# Import OpLog
op_log_path = os.path.join(model_dir, "metadata/tfprof_log")
op_log = tfprof_log_pb2.OpLog()
if gfile.Exists(op_log_path):
with gfile.GFile(op_log_path, "rb") as file:
op_log.MergeFromString(file.read())
print("Loaded OpLog from {}".format(op_log_path))
else:
print("OpLog does not exist a {}. Skipping.".format(op_log_path))
return run_meta, graph, op_log
def merge_default_with_oplog(graph, op_log=None, run_meta=None):
"""Monkeypatch. There currently is a bug in tfprof_logger that
prevents it from being used with Python 3. So we override the method
manually until the fix comes in.
"""
tmp_op_log = tfprof_log_pb2.OpLog()
# pylint: disable=W0212
logged_ops = tfprof_logger._get_logged_ops(graph, run_meta)
if not op_log:
tmp_op_log.log_entries.extend(logged_ops.values())
else:
all_ops = dict()
for entry in op_log.log_entries:
all_ops[entry.name] = entry
for op_name, entry in six.iteritems(logged_ops):
if op_name in all_ops:
all_ops[op_name].types.extend(entry.types)
if entry.float_ops > 0 and all_ops[op_name].float_ops == 0:
all_ops[op_name].float_ops = entry.float_ops
else:
all_ops[op_name] = entry
tmp_op_log.log_entries.extend(all_ops.values())
return tmp_op_log
## profile the above code
X_test, test_locs = build_toy_dataset(500)
from tensorflow.tools.tfprof import tfprof_log_pb2
run_metadata = tf.RunMetadata()
loss_value = sess.run(
inference.loss,
feed_dict={
X_ph: X_test,
y_ph: test_locs
},
options=tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE),
run_metadata=run_metadata)
op_log = tfprof_log_pb2.OpLog()
opts = tf.contrib.tfprof.model_analyzer.PRINT_ALL_TIMING_MEMORY.copy()
# opts['output'] = 'timeline:outfile=/home/mayank/temp/test_tfprof'
tf.contrib.tfprof.tfprof_logger.write_op_log(tf.get_default_graph(),
log_dir="/tmp/log_dir",
op_log=op_log,
run_meta=run_metadata)
bb = tf.contrib.tfprof.model_analyzer.print_model_analysis(
tf.get_default_graph(),
run_meta=run_metadata,
op_log=op_log,
tfprof_cmd='code',
tfprof_options=opts)
