Repository: chainer/chainer

Issue:
Add remove_variable option to build_computational_graph
In current implementation, the computational graph is a bipartite graph consisting in `Variable`s and `Function`s. Sometimes we do not need `Variable`s in the visualization of the computational graphs (e.g. Generated graph is too large. We just need to know how `Function`s are composed). To satisfy this requirement, it would be nice to extract `Functions` only from the computational graph and visualize them.

One of the possible implementation is to add `remove_variable` option to `build_computational_graph`.



Discussion:

Resolving patch:
diff --git a/chainer/computational_graph.py b/chainer/computational_graph.py
--- a/chainer/computational_graph.py
+++ b/chainer/computational_graph.py
@@ -3,6 +3,9 @@
 from chainer import function
 from chainer import variable
 
+_var_style = {'shape': 'octagon', 'fillcolor': '#E0E0E0', 'style': 'filled'}
+_func_style = {'shape': 'record', 'fillcolor': '#6495ED', 'style': 'filled'}
+
 
 class DotNode(object):
     """Node of the computational graph, with utilities for dot language.
@@ -13,15 +16,20 @@ class DotNode(object):
     Args:
         node: :class: `Variable` object or :class: `Function` object.
         attribute (dict): Attributes for the node.
+        show_name (bool): If `True`, the `name` attribute of the node is added
+            to the label. Default is `True`.
 
     """
 
-    def __init__(self, node, attribute=None):
+    def __init__(self, node, attribute=None, show_name=True):
         assert isinstance(node, (variable.Variable, function.Function))
         self.node = node
         self.id_ = id(node)
         self.attribute = {'label': node.label}
         if isinstance(node, variable.Variable):
+            if show_name and hasattr(node, 'name') and node.name is not None:
+                self.attribute['label'] = '{}: {}'.format(
+                    node.name, self.attribute['label'])
             self.attribute.update({'shape': 'oval'})
         else:
             self.attribute.update({'shape': 'box'})
@@ -59,11 +67,26 @@ class ComputationalGraph(object):
         rankdir (str): Direction of the graph that must be
             TB (top to bottom), BT (bottom to top), LR (left to right)
             or RL (right to left).
+        remove_variable (bool): If `True`, :class:`~chainer.Variable`s are
+            removed from the resulting computational graph. Only
+            :class:`~chainer.Function`s are shown in the output.
+        show_name (bool): If `True`, the `name` attribute of each node is
+            added to the label of the node. Default is `True`.
+
+    .. note::
+
+        The default behavior of :class:`~chainer.ComputationalGraph` has been
+        changed from v1.23.0, so that it ouputs the richest representation of
+        a graph as default, namely, styles are set and names of functions and
+        variables are shown. To reproduce the same result as previous versions
+        (<= v1.22.0), please specify `variable_style=None`,
+        `function_style=None`, and `show_name=False` explicitly.
 
     """
 
-    def __init__(self, nodes, edges, variable_style=None, function_style=None,
-                 rankdir='TB'):
+    def __init__(self, nodes, edges, variable_style=_var_style,
+                 function_style=_func_style, rankdir='TB',
+                 remove_variable=False, show_name=True):
         self.nodes = nodes
         self.edges = edges
         self.variable_style = variable_style
@@ -71,6 +94,8 @@ def __init__(self, nodes, edges, variable_style=None, function_style=None,
         if rankdir not in ('TB', 'BT', 'LR', 'RL'):
             raise ValueError('rankdir must be in TB, BT, LR or RL.')
         self.rankdir = rankdir
+        self.remove_variable = remove_variable
+        self.show_name = show_name
 
     def _to_dot(self):
         """Converts graph in dot format.
@@ -81,12 +106,20 @@ def _to_dot(self):
 
         """
         ret = 'digraph graphname{rankdir=%s;' % self.rankdir
+
+        if self.remove_variable:
+            self.nodes, self.edges = _skip_variable(self.nodes, self.edges)
+
         for node in self.nodes:
             assert isinstance(node, (variable.Variable, function.Function))
             if isinstance(node, variable.Variable):
-                ret += DotNode(node, self.variable_style).label
+                if not self.remove_variable:
+                    ret += DotNode(
+                        node, self.variable_style, self.show_name).label
             else:
-                ret += DotNode(node, self.function_style).label
+                ret += DotNode(node, self.function_style, self.show_name).label
+
+        drawn_edges = []
         for edge in self.edges:
             head, tail = edge
             if (isinstance(head, variable.Variable) and
@@ -98,11 +131,19 @@ def _to_dot(self):
                 head_attr = self.function_style
                 tail_attr = self.variable_style
             else:
-                raise TypeError(
-                    'head and tail should be the set of Variable and Function')
-            head_node = DotNode(head, head_attr)
-            tail_node = DotNode(tail, tail_attr)
-            ret += "%s -> %s;" % (head_node.id_, tail_node.id_)
+                if not self.remove_variable:
+                    raise TypeError('head and tail should be the set of '
+                                    'Variable and Function')
+                else:
+                    head_attr = self.function_style
+                    tail_attr = self.function_style
+            head_node = DotNode(head, head_attr, self.show_name)
+            tail_node = DotNode(tail, tail_attr, self.show_name)
+            edge = (head_node.id_, tail_node.id_)
+            if edge in drawn_edges:
+                continue
+            ret += "%s -> %s;" % edge
+            drawn_edges.append(edge)
         ret += "}"
         return ret
 
@@ -123,9 +164,34 @@ def dump(self, format='dot'):
             NotImplementedError('Currently, only dot format is supported.')
 
 
-def build_computational_graph(outputs, remove_split=True,
-                              variable_style=None, function_style=None,
-                              rankdir='TB'):
+def _skip_variable(nodes, edges):
+    func_edges = []
+    for edge_i, edge in enumerate(edges):
+        head, tail = edge
+        if isinstance(head, variable.Variable):
+            if head.creator is not None:
+                head = head.creator
+            else:
+                continue
+        if isinstance(tail, variable.Variable):
+            for node in nodes:
+                if isinstance(node, function.Function):
+                    for input_var in node.inputs:
+                        if input_var is tail:
+                            tail = node
+                            break
+                    if isinstance(tail, function.Function):
+                        break
+            else:
+                continue
+        func_edges.append((head, tail))
+    return nodes, func_edges
+
+
+def build_computational_graph(
+        outputs, remove_split=True, variable_style=_var_style,
+        function_style=_func_style, rankdir='TB', remove_variable=False,
+        show_name=True):
     """Builds a graph of functions and variables backward-reachable from outputs.
 
     Args:
@@ -140,6 +206,11 @@ def build_computational_graph(outputs, remove_split=True,
         rankdir (str): Direction of the graph that must be
             TB (top to bottom), BT (bottom to top), LR (left to right)
             or RL (right to left).
+        remove_variable (bool): If `True`, :class:`~chainer.Variable`s are
+            removed from the resulting computational graph. Only
+            :class:`~chainer.Function`s are shown in the output.
+        show_name (bool): If `True`, the `name` attribute of each node is
+            added to the label of the node. Default is `True`.
 
     Returns:
         ComputationalGraph: A graph consisting of nodes and edges that
@@ -166,6 +237,15 @@ def build_computational_graph(outputs, remove_split=True,
 
         See :class:`TestGraphBuilder` for details.
 
+    .. note::
+
+        The default behavior of :class:`~chainer.ComputationalGraph` has been
+        changed from v1.23.0, so that it ouputs the richest representation of
+        a graph as default, namely, styles are set and names of functions and
+        variables are shown. To reproduce the same result as previous versions
+        (<= v1.22.0), please specify `variable_style=None`,
+        `function_style=None`, and `show_name=False` explicitly.
+
     """
     if not remove_split:
         raise ValueError('remove_split=False is not supported anymore')
@@ -211,5 +291,6 @@ def add_cand(cand):
                     seen_edges.add((input_, cand))
                     nodes.add(HashableObject(input_))
                     nodes.add(HashableObject(cand))
-    return ComputationalGraph(list(i.v for i in nodes), list(seen_edges),
-                              variable_style, function_style, rankdir)
+    return ComputationalGraph(
+        list(i.v for i in nodes), list(seen_edges), variable_style,
+        function_style, rankdir, remove_variable, show_name)


Test patch:
diff --git a/tests/chainer_tests/test_computational_graph.py b/tests/chainer_tests/test_computational_graph.py
--- a/tests/chainer_tests/test_computational_graph.py
+++ b/tests/chainer_tests/test_computational_graph.py
@@ -230,4 +230,18 @@ def test_randir_invalid(self):
                           c.build_computational_graph, (self.y,), rankdir='TL')
 
 
+class TestGraphBuilderRemoveVariable(unittest.TestCase):
+
+    def setUp(self):
+        self.x1 = variable.Variable(np.zeros((1, 2)).astype('f'))
+        self.x2 = variable.Variable(np.zeros((1, 2)).astype('f'))
+        self.y = self.x1 + self.x2
+        self.f = self.y.creator
+        self.g = c.build_computational_graph((self.y,), remove_variable=True)
+
+    def test_remove_variable(self):
+        self.assertIn(self.f.label, self.g.dump())
+        self.assertNotIn(str(id(self.x1)), self.g.dump())
+        self.assertNotIn(str(id(self.x2)), self.g.dump())
+
 testing.run_module(__name__, __file__)