Merge pull request #8 from abergeron/shared_pickle

theano/sandbox/cuda/__init__.py

@@ -2,6 +2,7 @@ import atexit, logging, os, stat, sys
 from theano.compile import optdb
 from theano.gof.cmodule import get_lib_extension
 from theano.configparser import config, AddConfigVar, StrParam
+from theano.tensor.sharedvar import load_shared_variable
 import nvcc_compiler
 
 _logger_name = 'theano.sandbox.cuda'
@@ -275,11 +276,25 @@ def handle_shared_float32(tf):
     """
     if tf:
         import theano.compile
+        import copy_reg
         theano.compile.shared_constructor(float32_shared_constructor)
-
+        # this is a bit of hackery to make the shared variables load
+        # with the proper type.
+        copy_reg.pickle(theano.tensor.basic.TensorVariable, 
+                        reduce_tensor_variable, 
+                        load_shared_variable)
     else:
         raise NotImplementedError('removing our handler')
 
+def reduce_tensor_variable(var):
+    if isinstance(var.owner.op, HostFromGpu) and len(var.owner.inputs) == 1 \
+            and isinstance(var.owner.inputs[0], CudaNdarraySharedVariable):
+        return load_shared_variable, (var.owner.inputs[0].get_value(),)
+    else:
+        # this will make protocol 2 a little bit less efficient
+        # but there is no way around it.
+        return var.__reduce__()
+
 if config.device.startswith('gpu'):
     use(device=config.device, force=config.force_device)
 elif config.init_gpu_device:

theano/sandbox/cuda/tests/test_var.py

@@ -4,10 +4,29 @@ from nose.plugins.skip import SkipTest
 from theano.sandbox.cuda.var import float32_shared_constructor as f32sc
 from theano.sandbox.cuda import CudaNdarrayType, cuda_available
 
+import theano
+
 # Skip test if cuda_ndarray is not available.
 if cuda_available == False:
     raise SkipTest('Optional package cuda disabled')
 
+def test_shared_pickle():
+    import pickle
+    picklestring = "ctheano.tensor.sharedvar\nload_shared_variable\np0\n(cnumpy.core.multiarray\n_reconstruct\np1\n(cnumpy\nndarray\np2\n(I0\ntp3\nS'b'\np4\ntp5\nRp6\n(I1\n(I2\ntp7\ncnumpy\ndtype\np8\n(S'f4'\np9\nI0\nI1\ntp10\nRp11\n(I3\nS'<'\np12\nNNNI-1\nI-1\nI0\ntp13\nbI00\nS'\\x00\\x00\\x80?\\x00\\x00\\x00@'\np14\ntp15\nbtp16\nRp17\n."
+    g = pickle.loads(picklestring)
+    v = numpy.array([1.0, 2.0], dtype='float32')
+
+    # This test will always be on the GPU
+    assert isinstance(g, theano.tensor.basic.TensorVariable)
+    assert isinstance(g.owner, theano.gof.graph.Apply)
+    assert isinstance(g.owner.op, theano.sandbox.cuda.HostFromGpu)
+    assert isinstance(g.owner.inputs[0], CudaNdarrayType.SharedVariable)
+    assert (g.owner.inputs[0].get_value() == v).all()
+
+    # Make sure it saves the same way (so that the tests before are not bogus)
+    s = theano.tensor.as_tensor_variable(theano.shared(v))
+    assert pickle.dumps(s) == picklestring
+
 def test_float32_shared_constructor():
 
     npy_row = numpy.zeros((1,10), dtype='float32')

theano/tensor/sharedvar.py

@@ -5,9 +5,17 @@ from basic import TensorType, _tensor_py_operators, autocast_int, autocast_float
 from theano.compile import shared_constructor, SharedVariable
 from theano import config
 
+def load_shared_variable(val):
+    return theano.tensor.as_tensor_variable(theano.shared(val))
+
 # _tensor_py_operators is first to have its version of __{gt,ge,lt,le}__
 class TensorSharedVariable(_tensor_py_operators, SharedVariable):
-    pass
+    def __reduce_ex__(self, proto):
+        # This is for loading on the GPU if present.
+        if self.dtype == 'float32':
+            return load_shared_variable, (self.get_value(),)
+        else:
+            return super(TensorSharedVariable, self).__reduce_ex__(proto)
 
 @shared_constructor
 def tensor_constructor(value, name=None, strict=False, allow_downcast=None, borrow=False, broadcastable=None):

theano/tensor/tests/test_sharedvar.py

@@ -637,3 +637,25 @@ test_shared_options=makeSharedTester(
     cast_value_ = numpy.asarray,
     op_by_matrix_ = False,
     name='test_shared_options')
+
+
+def test_shared_pickle():
+    import pickle
+    picklestring = "ctheano.tensor.sharedvar\nload_shared_variable\np0\n(cnumpy.core.multiarray\n_reconstruct\np1\n(cnumpy\nndarray\np2\n(I0\ntp3\nS'b'\np4\ntp5\nRp6\n(I1\n(I2\ntp7\ncnumpy\ndtype\np8\n(S'f4'\np9\nI0\nI1\ntp10\nRp11\n(I3\nS'<'\np12\nNNNI-1\nI-1\nI0\ntp13\nbI00\nS'\\x00\\x00\\x80?\\x00\\x00\\x00@'\np14\ntp15\nbtp16\nRp17\n."
+    g = pickle.loads(picklestring)
+    v = numpy.array([1.0, 2.0], dtype='float32')
+
+    assert g.type == theano.tensor.fvector
+    
+    if theano.config.device.startswith('cpu'):
+        assert isinstance(g, theano.tensor.sharedvar.TensorSharedVariable)
+        assert (g.get_value() == v).all()
+    if theano.config.device.startswith('gpu'):
+        assert isinstance(g, theano.tensor.basic.TensorVariable)
+        # we don't go digging deeper because we don't want to 
+        # import theano.sandbox.cuda.
+        # some other tests are there.
+
+    # Make sure it saves the same way (so that the tests before are not bogus)
+    s = theano.tensor.as_tensor_variable(theano.shared(v))
+    assert pickle.dumps(s) == picklestring