def testRegisteredNode(self):
   graph = ops.Graph()
   node = ops._NodeDef("a", "an_a")
   weight_params = ops.get_stats_for_node_def(graph, node, "weight_parameters")
   self.assertEqual(10, weight_params.value)
   flops = ops.get_stats_for_node_def(graph, node, "flops")
   self.assertEqual(20, flops.value)
   missing_stat = ops.get_stats_for_node_def(graph, node, "missing_stat")
   self.assertEqual(None, missing_stat.value)

def calculate_graph_metrics(graph_def, statistic_types, input_layer,
                            input_shape_override, batch_size):
  """Looks at the performance statistics of all nodes in the graph."""
  _ = tf.import_graph_def(graph_def, name="")
  total_stats = {}
  node_stats = {}
  for statistic_type in statistic_types:
    total_stats[statistic_type] = ops.OpStats(statistic_type)
    node_stats[statistic_type] = {}
  # Make sure we get pretty-printed numbers with separators.
  locale.setlocale(locale.LC_ALL, "")
  with tf.Session() as sess:
    input_tensor = sess.graph.get_tensor_by_name(input_layer)
    input_shape_tensor = input_tensor.get_shape()
    if input_shape_tensor:
      input_shape = input_shape_tensor.as_list()
    else:
      input_shape = None
    if input_shape_override:
      input_shape = input_shape_override
    input_shape[0] = batch_size
    input_tensor.set_shape(input_shape)
    for node in graph_def.node:
      # Ensure that the updated input shape has been fully-propagated before we
      # ask for the statistics, since they may depend on the output size.
      op = sess.graph.get_operation_by_name(node.name)
      ops.set_shapes_for_outputs(op)
      for statistic_type in statistic_types:
        current_stats = ops.get_stats_for_node_def(sess.graph, node,
                                                   statistic_type)
        node_stats[statistic_type][node.name] = current_stats
        total_stats[statistic_type] += current_stats
  return total_stats, node_stats

def main(unused_args):
    if not tf.gfile.Exists(FLAGS.graph):
        print("Input graph file '" + FLAGS.graph + "' does not exist!")
        return -1

    graph_def = graph_pb2.GraphDef()
    with open(FLAGS.graph, "rb") as f:
        if FLAGS.input_binary:
            graph_def.ParseFromString(f.read())
        else:
            text_format.Merge(f.read(), graph_def)
    _ = tf.import_graph_def(graph_def, name="")

    statistic_types = FLAGS.statistics.split(",")
    total_stats = {}
    for statistic_type in statistic_types:
        total_stats[statistic_type] = ops.OpStats(statistic_type)
    with tf.Session() as sess:
        input_tensor = sess.graph.get_tensor_by_name(FLAGS.input_layer)
        input_shape = input_tensor.get_shape()
        input_shape = [FLAGS.batch_size, input_shape[1], input_shape[2], input_shape[3]]
        input_tensor.set_shape(input_shape)
        for node in graph_def.node:
            for statistic_type in statistic_types:
                node_stats = ops.get_stats_for_node_def(sess.graph, node, statistic_type)
                total_stats[statistic_type] += node_stats
    # Make sure we get pretty-printed numbers with separators.
    locale.setlocale(locale.LC_ALL, "")
    for statistic_type in statistic_types:
        value = total_stats[statistic_type].value
        if value is None:
            friendly_value = "None"
        else:
            friendly_value = locale.format("%d", value, grouping=True)
        print("%s=%s" % (statistic_type, friendly_value))

def _get_logged_ops(graph):
  """Extract trainable model parameters and FLOPs for ops from a Graph.

  Args:
    graph: tf.Graph.
  Returns:
    logged_ops: dict mapping from op_name to OpLogEntry.
  """
  logged_ops = {}

  graph_def = graph.as_graph_def()
  for node in graph_def.node:
    try:
      stats = ops.get_stats_for_node_def(graph, node, REGISTERED_FLOP_STATS)
    except ValueError:
      # Catch Exception When shape is incomplete. Skip it.
      stats = None

    if not stats or not stats.value:
      continue
    if node.name not in logged_ops:
      entry = tfprof_log_pb2.OpLogEntry()
      entry.name = node.name
      entry.float_ops = stats.value
      logged_ops[entry.name] = entry

  for v in graph.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES):
    if v.op.name not in logged_ops:
      entry = tfprof_log_pb2.OpLogEntry()
      entry.name = v.op.name
      entry.types.append(TRAINABLE_VARIABLES)
      logged_ops[entry.name] = entry
    else:
      logged_ops[v.op.name].types.append(TRAINABLE_VARIABLES)
  return logged_ops

 def testTransposedStatistics(self):
   a = variables.Variable(random_ops.random_normal([16, 25]))
   b = variables.Variable(random_ops.random_normal([16, 9]))
   math_ops.matmul(a, b, transpose_a=True)
   g = ops.get_default_graph()
   for op in g.get_operations():
     flops = ops.get_stats_for_node_def(g, op.node_def, "flops").value
     if op.name == "MatMul":
       self.assertEqual(7200, flops)

def _get_logged_ops(graph, run_meta=None, add_trace=True):
  """Extract trainable model parameters and FLOPs for ops from a Graph.

  Args:
    graph: tf.Graph.
    run_meta: RunMetadata proto used to complete shape information.
    add_trace: Whether to add op trace information.
  Returns:
    logged_ops: dict mapping from op_name to OpLogEntry.
  """
  if run_meta:
    graph = _fill_missing_graph_shape(graph, run_meta)

  op_missing_shape = 0
  logged_ops = {}
  for op in graph.get_operations():
    try:
      stats = ops.get_stats_for_node_def(
          graph, op.node_def, REGISTERED_FLOP_STATS)
    except ValueError:
      # Catch Exception When shape is incomplete. Skip it.
      op_missing_shape += 1
      stats = None

    entry = tfprof_log_pb2.OpLogEntry()
    entry.name = op.name
    add_entry = False
    if stats and stats.value:
      entry.float_ops = int(stats.value)
      add_entry = True

    if add_trace:
      for tb in op.traceback:
        trace = entry.code_def.traces.add()
        trace.file = tb[0] if tb[0] else 'none'
        trace.lineno = tb[1] if tb[1] else -1
        trace.function = tb[2] if tb[2] else 'none'
        trace.line = tb[3] if tb[3] else 'none'
      add_entry = True

    if add_entry:
      logged_ops[entry.name] = entry

  for v in graph.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES):
    if v.op.name not in logged_ops:
      entry = tfprof_log_pb2.OpLogEntry()
      entry.name = v.op.name
      entry.types.append(TRAINABLE_VARIABLES)
      logged_ops[entry.name] = entry
    else:
      logged_ops[v.op.name].types.append(TRAINABLE_VARIABLES)
  if op_missing_shape > 0 and not run_meta:
    sys.stderr.write('%d ops no flops stats due to incomplete shapes. '
                     'Consider passing run_meta to use run_time shapes.\n' %
                     op_missing_shape)
  return logged_ops

 def testSimpleStatistics(self):
   g = ops.Graph()
   with g.as_default():
     a = variables.Variable(random_ops.random_normal([25, 16]))
     b = variables.Variable(random_ops.random_normal([16, 9]))
     math_ops.matmul(a, b)
     for op in g.get_operations():
       flops = ops.get_stats_for_node_def(g, op.node_def, "flops").value
       if op.name == "MatMul":
         self.assertEqual(7200, flops)

 def testTransposedStatistics(self):
   g = tf.Graph()
   with g.as_default():
     a = tf.Variable(tf.random_normal([16, 25]))
     b = tf.Variable(tf.random_normal([16, 9]))
     tf.matmul(a, b, transpose_a=True)
     for op in g.get_operations():
       flops = ops.get_stats_for_node_def(g, op.node_def, "flops").value
       if op.name == "MatMul":
         self.assertEqual(7200, flops)

def _flops(op):
  """Get the number of flops of a convolution, from the ops stats registry.

  Args:
    op: A tf.Operation object.

  Returns:
    The number os flops needed to evaluate conv_op.
  """
  return (ops.get_stats_for_node_def(tf.get_default_graph(), op.node_def,
                                     'flops').value)

def _get_logged_ops(graph, run_meta=None):
  """Extract trainable model parameters and FLOPs for ops from a Graph.

  Args:
    graph: tf.Graph.
    run_meta: RunMetadata proto used to complete shape information.
  Returns:
    logged_ops: dict mapping from op_name to OpLogEntry.
  """
  if run_meta:
    graph = _fill_missing_graph_shape(graph, run_meta)

  op_missing_shape = 0
  logged_ops = {}
  graph_def = graph.as_graph_def()
  for node in graph_def.node:
    try:
      stats = ops.get_stats_for_node_def(graph, node, REGISTERED_FLOP_STATS)
    except ValueError:
      # Catch Exception When shape is incomplete. Skip it.
      op_missing_shape += 1
      stats = None

    if not stats or not stats.value:
      continue
    if node.name not in logged_ops:
      entry = tfprof_log_pb2.OpLogEntry()
      entry.name = node.name
      entry.float_ops = int(stats.value)
      logged_ops[entry.name] = entry

  for v in graph.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES):
    if v.op.name not in logged_ops:
      entry = tfprof_log_pb2.OpLogEntry()
      entry.name = v.op.name
      entry.types.append(TRAINABLE_VARIABLES)
      logged_ops[entry.name] = entry
    else:
      logged_ops[v.op.name].types.append(TRAINABLE_VARIABLES)
  if op_missing_shape > 0 and not run_meta:
    sys.stderr.write(
        '%d ops no flops stats due to incomplete shapes. '
        'Consider passing run_meta to use run_time shapes.\n' %
        op_missing_shape)
  return logged_ops

 def testUnregisteredNode(self):
   graph = ops.Graph()
   node = ops._NodeDef("b", "a_b")
   weight_params = ops.get_stats_for_node_def(graph, node, "weight_params")
   self.assertEqual(None, weight_params.value)

def _get_logged_ops(graph, run_meta=None, add_trace=True,
                    add_trainable_var=True):
  """Extract trainable model parameters and FLOPs for ops from a Graph.

  Args:
    graph: tf.Graph.
    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:
    logged_ops: dict mapping from op_name to OpLogEntry.
    string_to_id: dict mapping from string to id.
  """
  if run_meta:
    graph = _fill_missing_graph_shape(graph, run_meta)

  op_missing_shape = 0
  logged_ops = {}
  string_to_id = dict()
  string_to_id['none'] = len(string_to_id)
  # TODO(xpan): Work with Profiler more efficiently.
  for op in graph.get_operations():
    try:
      stats = ops.get_stats_for_node_def(
          graph, op.node_def, REGISTERED_FLOP_STATS)
    except ValueError:
      # Catch Exception When shape is incomplete. Skip it.
      op_missing_shape += 1
      stats = None

    entry = tfprof_log_pb2.OpLogEntry()
    entry.name = op.name
    add_entry = False
    if stats and stats.value:
      entry.float_ops = int(stats.value)
      add_entry = True

    if add_trace:
      for tb in op.traceback_with_start_lines:
        trace = entry.code_def.traces.add()
        trace.file_id = _str_id(tb[0], string_to_id) if tb[0] else 0
        trace.lineno = tb[1] if tb[1] else -1
        trace.function_id = _str_id(tb[2], string_to_id) if tb[2] else 0
        trace.line_id = _str_id(tb[3], string_to_id) if tb[3] else 0
        trace.func_start_line = tb[4] if tb[4] else -1
      add_entry = True

    if add_entry:
      logged_ops[entry.name] = entry

  if add_trainable_var:
    for v in graph.get_collection(ops.GraphKeys.TRAINABLE_VARIABLES):
      if v.op.name not in logged_ops:
        entry = tfprof_log_pb2.OpLogEntry()
        entry.name = v.op.name
        entry.types.append(TRAINABLE_VARIABLES)
        logged_ops[entry.name] = entry
      else:
        logged_ops[v.op.name].types.append(TRAINABLE_VARIABLES)

  if op_missing_shape > 0 and not run_meta:
    sys.stderr.write('%d ops no flops stats due to incomplete shapes.\n' %
                     op_missing_shape)
  return logged_ops, string_to_id

def calculate_graph_metrics(graph_def, statistic_types, input_layer,
                            input_shape_override, batch_size):
    """Looks at the performance statistics of all nodes in the graph.

    Parameters
    ----------
    graph_def : TYPE
        Description
    statistic_types : TYPE
        Description
    input_layer : TYPE
        Description
    input_shape_override : TYPE
        Description
    batch_size : TYPE
        Description

    Returns
    -------
    TYPE
        Description

    Raises
    ------
    ValueError
        Description
    """
    tf.import_graph_def(graph_def, name="")
    total_stats = {}
    node_stats = {}
    for statistic_type in statistic_types:
        total_stats[statistic_type] = ops.OpStats(statistic_type)
        node_stats[statistic_type] = {}
    # Make sure we get pretty-printed numbers with separators.
    locale.setlocale(locale.LC_ALL, "")
    with tf.Session() as sess:
        input_tensor = sess.graph.get_tensor_by_name(input_layer)
        input_shape_tensor = input_tensor.get_shape()
        if input_shape_tensor:
            input_shape = input_shape_tensor.as_list()
        else:
            input_shape = None
        if input_shape_override:
            input_shape = input_shape_override
        if input_shape is None:
            raise ValueError("""No input shape was provided on the command line,"""
                             """ and the input op itself had no default shape, so"""
                             """ shape inference couldn't be performed. This is"""
                             """ required for metrics calculations.""")
        input_shape[0] = batch_size
        input_tensor.set_shape(input_shape)
        for node in graph_def.node:
            # Ensure that the updated input shape has been fully-propagated before we
            # ask for the statistics, since they may depend on the output size.
            op = sess.graph.get_operation_by_name(node.name)
            ops.set_shapes_for_outputs(op)
            for statistic_type in statistic_types:
                current_stats = ops.get_stats_for_node_def(sess.graph, node,
                                                           statistic_type)
                node_stats[statistic_type][node.name] = current_stats
                total_stats[statistic_type] += current_stats
    return total_stats, node_stats