Repository: chainer/chainer

Issue:
Can't pickle optimizer from chainer>=6.2.0
```python
import chainer
import pickle

lnk = chainer.links.Linear(1, 2)
opt = chainer.optimizers.Adam().setup(lnk)
print(pickle.dumps(opt))
```

The above code is fine for chainer<=6.1.0 but raises an error for >=6.2.0.

```python
Traceback (most recent call last):
  File "pickle_opt.py", line 6, in <module>
    print(pickle.dumps(opt))
AttributeError: Can't pickle local object 'Optimizer.setup.<locals>.OptimizerHookable'
```

I'm not sure if users can expect optimizers to be pickle-able, but this can be a problem e.g. when passing an optimizer to another process for multiprocessing.


Discussion:

Resolving patch:
diff --git a/chainer/optimizer.py b/chainer/optimizer.py
--- a/chainer/optimizer.py
+++ b/chainer/optimizer.py
@@ -497,6 +497,17 @@ def use_fp32_update(self, flag=True):
         self._use_fp32_update = flag
 
 
+class _OptimizerHookable(_Hookable):
+    def __init__(self, optimizer):
+        super(_OptimizerHookable, self).__init__(
+            invalid_timing_fallback=True)
+        self.optimizer = optimizer
+
+    def call_hook(self, hook, args):
+        assert args == ()
+        self.optimizer.call_hook(hook)
+
+
 class Optimizer(object):
     """Base class of all numerical optimizers.
 
@@ -571,18 +582,7 @@ def setup(self, link):
         self.t = 0
         self.epoch = 0
 
-        optimizer = self
-
-        class OptimizerHookable(_Hookable):
-            def __init__(self):
-                super(OptimizerHookable, self).__init__(
-                    invalid_timing_fallback=True)
-
-            def call_hook(self, hook, args):
-                assert args == ()
-                optimizer.call_hook(hook)
-
-        self._hookable = OptimizerHookable()
+        self._hookable = _OptimizerHookable(self)
         return self
 
     def update(self, lossfun=None, *args, **kwds):


Test patch:
diff --git a/tests/chainer_tests/optimizers_tests/test_optimizers.py b/tests/chainer_tests/optimizers_tests/test_optimizers.py
--- a/tests/chainer_tests/optimizers_tests/test_optimizers.py
+++ b/tests/chainer_tests/optimizers_tests/test_optimizers.py
@@ -1,3 +1,4 @@
+import pickle
 import unittest
 
 import numpy as np
@@ -212,6 +213,50 @@ def test_hooks_auto(self):
         self.assertNotEqual(h_pre.value, h_post.value)
 
 
+@_parameterize_optimizers
+class TestOptimizerPickable(unittest.TestCase):
+
+    def setUp(self):
+        self.target = SimpleChain()
+
+    def create(self, *args, **kwargs):
+        self.optimizer = self.optimizer_impl(*args, **kwargs)
+        self.optimizer.setup(self.target)
+
+    def get_hyperparam(self, name):
+        return getattr(self.target.w.update_rule.hyperparam, name)
+
+    def test_new_pickle(self):
+        self.create()
+        pickled_opt = pickle.dumps(self.optimizer)
+
+        x = chainer.Variable(np.array(5., dtype=np.float32))
+        self.optimizer.update(self.target, x)
+        w_post = np.copy(self.target.w.data)
+        # Pickle has saved a copy of the target
+        opt = pickle.loads(pickled_opt)
+        opt.update(opt.target, x)
+        pickled_w_post = np.copy(opt.target.w.data)
+
+        self.assertEqual(w_post, pickled_w_post)
+
+    def test_updated_pickle(self):
+        self.create()
+
+        x = chainer.Variable(np.array(5., dtype=np.float32))
+        self.optimizer.update(self.target, x)
+        pickled_opt = pickle.dumps(self.optimizer)
+
+        self.optimizer.update(self.target, x)
+        w_post = np.copy(self.target.w.data)
+        # Pickle has saved a copy of the target
+        opt = pickle.loads(pickled_opt)
+        opt.update(opt.target, x)
+        pickled_w_post = np.copy(opt.target.w.data)
+
+        self.assertEqual(w_post, pickled_w_post)
+
+
 @_parameterize_optimizers
 class TestOptimizerLossScaling(unittest.TestCase):