Repository: chainer/chainer

Issue:
Release training graph before validation
The observation dictionary currently holds the reported Variable object as is. If a forward computation code reports a Variable object, it holds the whole computational graph during the iteration. If an evaluation runs at that iteration, the computational graph is left on the memory during the evaluation, so that the maximum memory consumption becomes high. This problem can be fixed by not preserving the Variable object but instead holding the array value.

We have to think about how to retrieve the computational graph from the dump_graph extension. The current implementation requires an access to the whole computational graph via the observation dictionary. We need some workaround to give dump_graph extension an access to the computational graph.


Discussion:
One idea is using a global flag to configure if the reporter should clone reported variables and purge computational graphs from them. `dump_graph` is then modified to disable the purge at the first iteration, dumps the graph using the graph left at a reported variable at the second iteration, and then enable the purge again. Since an evaluator does not run in the second iteration in most cases, it should work.

I will work on this after #2145 is merged so that I can use the configuration tool to implement the global flag.

Resolving patch:
diff --git a/chainer/__init__.py b/chainer/__init__.py
--- a/chainer/__init__.py
+++ b/chainer/__init__.py
@@ -79,6 +79,8 @@ def get_function_hooks():
 
 global_config.debug = bool(int(os.environ.get('CHAINER_DEBUG', '0')))
 global_config.enable_backprop = True
+global_config.keep_graph_on_report = bool(int(
+    os.environ.get('CHAINER_KEEP_GRAPH_ON_REPORT', '0')))
 global_config.train = True
 global_config.type_check = bool(int(os.environ.get('CHAINER_TYPE_CHECK', '1')))
 global_config.use_cudnn = os.environ.get('CHAINER_USE_CUDNN', 'auto')
diff --git a/chainer/reporter.py b/chainer/reporter.py
--- a/chainer/reporter.py
+++ b/chainer/reporter.py
@@ -1,13 +1,21 @@
 import collections
 import contextlib
+import copy
 
 import numpy
 import six
 
+from chainer import configuration
 from chainer import cuda
 from chainer import variable
 
 
+def _copy_variable(value):
+    if isinstance(value, variable.Variable):
+        return copy.copy(value)
+    return value
+
+
 class Reporter(object):
 
     """Object to which observed values are reported.
@@ -126,6 +134,13 @@ def report(self, values, observer=None):
         The values are written with the key, prefixed by the name of the
         observer object if given.
 
+        .. note::
+           As of v2.0.0, if a value is of type :class:`~chainer.Variable`, the
+           variable is copied without preserving the computational graph and
+           the new variable object purged from the graph is stored to the
+           observer. This behavior can be changed by setting
+           ``chainer.config.keep_graph_on_report`` to ``True``.
+
         Args:
             values (dict): Dictionary of observed values.
             observer: Observer object. Its object ID is used to retrieve the
@@ -133,6 +148,9 @@ def report(self, values, observer=None):
                 name of the observed value.
 
         """
+        if not configuration.config.keep_graph_on_report:
+            values = {k: _copy_variable(v) for k, v in six.iteritems(values)}
+
         if observer is not None:
             observer_id = id(observer)
             if observer_id not in self._observer_names:
diff --git a/chainer/training/extensions/computational_graph.py b/chainer/training/extensions/computational_graph.py
--- a/chainer/training/extensions/computational_graph.py
+++ b/chainer/training/extensions/computational_graph.py
@@ -1,6 +1,7 @@
 import os
 
 from chainer import computational_graph
+from chainer import configuration
 from chainer.training import extension
 from chainer import variable
 
@@ -16,7 +17,33 @@ def dump_graph(root_name, out_name='cg.dot',
     This extension dumps a computational graph. The graph is output in DOT
     language.
 
-    It only dumps a graph at the first iteration by default.
+    It only dumps a graph at the first invocation.
+
+    .. note::
+       As of v2.0.0, the computational graph is not kept by default. This
+       extension changes this behavior until the first invocation. **It is
+       strongly recommended to use it with the default trigger setting.**
+
+       The detailed behavior of this extension since v2.0.0 is as follows.
+
+       1. In its initializer, it turns on the
+          ``chainer.config.keep_graph_on_report`` flag.
+       2. At the first iteration, it dumps the graph using the graph held by
+          the reported variable.
+       3. After dumping the graph, it turns off the flag (if it was originally
+          turned off) so that any variable reported afterward does not hold
+          a computational graph.
+
+       When the ``keep_graph_on_report`` flag is turned on, the computational
+       graph created by the updater is kept during the invocation of
+       extensions. It will cause an unnecessarily large memory consumption
+       when an extension also uses a large amount of memory, e.g.
+       :class:`~chainer.training.extensions.Evaluator`.
+
+       With the default setting, the ``dump_graph`` extension is called at the
+       first iteration. Since :class:`~chainer.training.extensions.Evaluator`
+       is not called at the first iteration in most cases, it does not cause
+       any memory problem.
 
     Args:
         root_name (str): Name of the root of the computational graph. The
@@ -41,20 +68,29 @@ def trigger(trainer):
     if function_style is None:
         function_style = _func_style
 
-    @extension.make_extension(trigger=trigger)
+    original_flag = [None]
+
+    def initializer(_):
+        original_flag[0] = configuration.config.keep_graph_on_report
+        configuration.config.keep_graph_on_report = True
+
+    @extension.make_extension(trigger=trigger, initializer=initializer)
     def dump_graph(trainer):
-        var = trainer.observation[root_name]
-        if not isinstance(var, variable.Variable):
-            raise TypeError('root value is not a Variable')
-        cg = computational_graph.build_computational_graph(
-            [var],
-            variable_style=variable_style,
-            function_style=function_style
-        ).dump()
-
-        out_path = os.path.join(trainer.out, out_name)
-        # TODO(beam2d): support outputting images by the dot command
-        with open(out_path, 'w') as f:
-            f.write(cg)
+        try:
+            var = trainer.observation[root_name]
+            if not isinstance(var, variable.Variable):
+                raise TypeError('root value is not a Variable')
+            cg = computational_graph.build_computational_graph(
+                [var],
+                variable_style=variable_style,
+                function_style=function_style
+            ).dump()
+
+            out_path = os.path.join(trainer.out, out_name)
+            # TODO(beam2d): support outputting images by the dot command
+            with open(out_path, 'w') as f:
+                f.write(cg)
+        finally:
+            configuration.config.keep_graph_on_report = original_flag[0]
 
     return dump_graph


Test patch:
diff --git a/tests/chainer_tests/test_reporter.py b/tests/chainer_tests/test_reporter.py
--- a/tests/chainer_tests/test_reporter.py
+++ b/tests/chainer_tests/test_reporter.py
@@ -1,9 +1,12 @@
+import contextlib
 import unittest
 
 import numpy
 
 import chainer
+from chainer import configuration
 from chainer import cuda
+from chainer import functions
 from chainer import testing
 from chainer.testing import attr
 
@@ -75,6 +78,47 @@ def test_report_without_observer(self):
         self.assertEqual(observation['x'], 1)
 
 
+class TestKeepGraphOnReportFlag(unittest.TestCase):
+
+    @contextlib.contextmanager
+    def _scope(self, flag):
+        # If flag is None, return the nop context.
+        # Otherwise, return the context in which
+        # keep_graph_on_report is set temporarily.
+        old = configuration.config.keep_graph_on_report
+        if flag is not None:
+            configuration.config.keep_graph_on_report = flag
+        try:
+            yield
+        finally:
+            configuration.config.keep_graph_on_report = old
+
+    def test_keep_graph_default(self):
+        x = chainer.Variable(numpy.array([1], numpy.float32))
+        y = functions.Sigmoid()(x)
+        reporter = chainer.Reporter()
+        with self._scope(None):
+            reporter.report({'y': y})
+        self.assertIsNone(reporter.observation['y'].creator)
+
+    def test_keep_graph(self):
+        x = chainer.Variable(numpy.array([1], numpy.float32))
+        y = functions.Sigmoid()(x)
+        reporter = chainer.Reporter()
+        with self._scope(True):
+            reporter.report({'y': y})
+        self.assertIsInstance(reporter.observation['y'].creator,
+                              functions.Sigmoid)
+
+    def test_not_keep_graph(self):
+        x = chainer.Variable(numpy.array([1], numpy.float32))
+        y = functions.Sigmoid()(x)
+        reporter = chainer.Reporter()
+        with self._scope(False):
+            reporter.report({'y': y})
+        self.assertIsNone(reporter.observation['y'].creator)
+
+
 class TestReport(unittest.TestCase):
 
     def test_report_without_reporter(self):
diff --git a/tests/chainer_tests/training_tests/extensions_tests/test_computational_graph.py b/tests/chainer_tests/training_tests/extensions_tests/test_computational_graph.py
new file mode 100644
--- /dev/null
+++ b/tests/chainer_tests/training_tests/extensions_tests/test_computational_graph.py
@@ -0,0 +1,103 @@
+import tempfile
+import os
+import shutil
+import unittest
+
+import numpy
+
+import chainer
+from chainer import configuration
+from chainer import functions
+from chainer import links
+from chainer import testing
+from chainer import training
+from chainer.training.extensions import computational_graph as c
+
+
+class Dataset(chainer.dataset.DatasetMixin):
+    def __init__(self, values):
+        self.values = values
+
+    def __len__(self):
+        return len(self.values)
+
+    def get_example(self, i):
+        return numpy.array([self.values[i]], numpy.float32), numpy.int32(i % 2)
+
+
+class Model(chainer.Link):
+    def __init__(self):
+        super(Model, self).__init__()
+        self.flag_history = []
+        self.l1 = links.Linear(2)
+        self.l2 = links.Linear(2)
+        self.i = 0
+
+    def __call__(self, x):
+        self.flag_history.append(
+            configuration.config.keep_graph_on_report)
+
+        h = self.l1(x)
+        if self.i == 0:
+            h = functions.Sigmoid()(h)
+        else:
+            h = functions.Tanh()(h)
+        h = self.l2(h)
+
+        self.i += 1
+        return h
+
+
+class TestGraphBuilderKeepGraphOnReport(unittest.TestCase):
+
+    def _run_test(self, tempdir, initial_flag):
+        n_data = 4
+        n_epochs = 3
+        outdir = os.path.join(tempdir, 'testresult')
+
+        # Prepare
+        model = Model()
+        classifier = links.Classifier(model)
+        optimizer = chainer.optimizers.Adam()
+        optimizer.setup(classifier)
+
+        dataset = Dataset([i for i in range(n_data)])
+        iterator = chainer.iterators.SerialIterator(dataset, 1, shuffle=False)
+        updater = training.StandardUpdater(iterator, optimizer)
+        trainer = training.Trainer(updater, (n_epochs, 'epoch'), out=outdir)
+
+        extension = c.dump_graph('main/loss', out_name='test.dot')
+        trainer.extend(extension)
+
+        # Run
+        with chainer.using_config('keep_graph_on_report', initial_flag):
+            trainer.run()
+
+        # Check flag history
+        self.assertEqual(model.flag_history,
+                         [True] + [initial_flag] * (n_data * n_epochs-1))
+
+        # Check the dumped graph
+        graph_path = os.path.join(outdir, 'test.dot')
+        with open(graph_path) as f:
+            graph_dot = f.read()
+
+        # Check that only the first iteration is dumped
+        self.assertIn('Sigmoid', graph_dot)
+        self.assertNotIn('Tanh', graph_dot)
+
+    def _check(self, initial_flag):
+        tempdir = tempfile.mkdtemp()
+        try:
+            self._run_test(tempdir, initial_flag)
+        finally:
+            shutil.rmtree(tempdir)
+
+    def test_keep_graph_on_report_flag_true(self):
+        self._check(True)
+
+    def test_keep_graph_on_report_flag_false(self):
+        self._check(False)
+
+
+testing.run_module(__name__, __file__)