GraphDef to DOT

graph_to_dot.py

def as_dot(g):
  names = dict()
  dot = "digraph {\n"
  for op_id in sorted(g._nodes_by_id):
    op = g._nodes_by_id[op_id]
    parts = op.name.split("/")
    n = names
    for p in parts:
      if p not in n:
        n[p] = dict()
      n = n[p]
    n[op_id] = op_id
    dot += str(op_id) + " [label=\"\\N : " + op.type + "\"];\n"
    for input_op in op.inputs:
      dot += str(input_op.op._id_value) + " -> " + str(op._id_value)
      dot += " [label=\"" + str(input_op.get_shape()) + "\"]"
      dot += ";\n"
    for input_op in op.control_inputs:
      dot += str(input_op._id_value) + " -> " + str(op._id_value) + " [style=\"dashed\"];\n"
  def print_dict(d, name):
    # if len(d) == 1: return str(d[d.keys()[0]]) + ";\n"
    dot = "subgraph cluster_" + name + " {\n"
    dot += "label=\"" + name + "\";\n"
    for k, v in d.iteritems():
      if isinstance(v, int): dot += str(v) + ";\n"
      if isinstance(v, dict): dot += print_dict(v, k)
    dot += "}\n"
    return dot
  dot += print_dict(names, "top")
  dot += "}"
  return dot