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)