Merge pull request #3926 from nouiz/gpuarray_pickle

theano/sandbox/gpuarray/tests/GpuArray.pkl

+ctheano.sandbox.gpuarray.type
+GpuArray_unpickler
+p0
+(cnumpy.core.multiarray
+_reconstruct
+p1
+(cnumpy
+ndarray
+p2
+(I0
+tp3
+S'b'
+p4
+tp5
+Rp6
+(I1
+(I1
+tp7
+cnumpy
+dtype
+p8
+(S'f4'
+p9
+I0
+I1
+tp10
+Rp11
+(I3
+S'<'
+p12
+NNNI-1
+I-1
+I0
+tp13
+bI00
+S'\x00\x00(\xc2'
+p14
+tp15
+bNtp16
+Rp17
+.
\ No newline at end of file

theano/sandbox/gpuarray/tests/test_pickle.py

+"""Some pickle test when pygpu isn't there. The test when pygpu is
+available are in test_type.py.
+
+This is needed as we skip all the test file when pygpu isn't there in
+regular test file.
+
+"""
+import os
+
+from nose.plugins.skip import SkipTest
+from nose.tools import assert_raises
+import numpy
+
+import theano.sandbox.gpuarray
+from theano.compat import PY3
+from theano import config
+from theano.misc.pkl_utils import CompatUnpickler
+
+if not theano.sandbox.gpuarray.pygpu_activated:
+    try:
+        import pygpu
+    except ImportError:
+        pygpu = None
+    import theano.sandbox.cuda as cuda_ndarray
+    if pygpu and cuda_ndarray.cuda_available:
+        cuda_ndarray.use('gpu', default_to_move_computation_to_gpu=False,
+                         move_shared_float32_to_gpu=False,
+                         enable_cuda=False)
+        theano.sandbox.gpuarray.init_dev('cuda')
+
+from .. import pygpu_activated  # noqa
+
+
+def test_unpickle_gpuarray_as_numpy_ndarray_flag1():
+    """Only test when pygpu isn't
+    available. test_unpickle_gpuarray_as_numpy_ndarray_flag0 in
+    test_type.py test it when pygpu is there.
+
+    """
+    if pygpu_activated:
+        raise SkipTest("pygpu disabled")
+    oldflag = config.experimental.unpickle_gpu_on_cpu
+    config.experimental.unpickle_gpu_on_cpu = False
+
+    try:
+        testfile_dir = os.path.dirname(os.path.realpath(__file__))
+        fname = 'GpuArray.pkl'
+
+        with open(os.path.join(testfile_dir, fname), 'rb') as fp:
+            if PY3:
+                u = CompatUnpickler(fp, encoding="latin1")
+            else:
+                u = CompatUnpickler(fp)
+            assert_raises(ImportError, u.load)
+    finally:
+        config.experimental.unpickle_gpu_on_cpu = oldflag
+
+
+def test_unpickle_gpuarray_as_numpy_ndarray_flag2():
+    oldflag = config.experimental.unpickle_gpu_on_cpu
+    config.experimental.unpickle_gpu_on_cpu = True
+
+    try:
+        testfile_dir = os.path.dirname(os.path.realpath(__file__))
+        fname = 'GpuArray.pkl'
+
+        with open(os.path.join(testfile_dir, fname), 'rb') as fp:
+            if PY3:
+                u = CompatUnpickler(fp, encoding="latin1")
+            else:
+                u = CompatUnpickler(fp)
+            mat = u.load()
+
+        assert isinstance(mat, numpy.ndarray)
+        assert mat[0] == -42.0
+
+    finally:
+        config.experimental.unpickle_gpu_on_cpu = oldflag

theano/sandbox/gpuarray/tests/test_type.py

+import os
+
 import numpy
 
 import theano
+from theano.compat import PY3
+from theano import config
 from theano.compile import DeepCopyOp
+from theano.misc.pkl_utils import CompatUnpickler
 
 from .config import test_ctx_name
 from .test_basic_ops import rand_gpuarray
+from ..type import GpuArrayType, gpuarray_shared_constructor
 
-from ..type import GpuArrayType
+import pygpu
 
 
 def test_deep_copy():
@@ -41,5 +47,31 @@ def test_specify_shape():
 
 
 def test_filter_float():
-    s = theano.shared(numpy.array(0.0, dtype='float32'), target=test_ctx_name)
+    theano.compile.shared_constructor(gpuarray_shared_constructor)
+    try:
+        s = theano.shared(numpy.array(0.0, dtype='float32'),
+                          target=test_ctx_name)
         theano.function([], updates=[(s, 0.0)])
+    finally:
+        del theano.compile.sharedvalue.shared.constructors[-1]
+
+
+def test_unpickle_gpuarray_as_numpy_ndarray_flag0():
+    """ Test when pygpu isn't there for unpickle are in test_pickle.py"""
+    oldflag = config.experimental.unpickle_gpu_on_cpu
+    config.experimental.unpickle_gpu_on_cpu = False
+
+    try:
+        testfile_dir = os.path.dirname(os.path.realpath(__file__))
+        fname = 'GpuArray.pkl'
+
+        with open(os.path.join(testfile_dir, fname), 'rb') as fp:
+            if PY3:
+                u = CompatUnpickler(fp, encoding="latin1")
+            else:
+                u = CompatUnpickler(fp)
+            mat = u.load()
+            assert isinstance(mat, pygpu.gpuarray.GpuArray)
+            assert numpy.asarray(mat)[0] == -42.0
+    finally:
+        config.experimental.unpickle_gpu_on_cpu = oldflag

theano/sandbox/gpuarray/type.py

 import numpy
+import six.moves.copyreg as copyreg
+from six import iteritems
+import warnings
 
 import theano
+from theano.tensor.type import TensorType
 from theano.tensor.var import _tensor_py_operators
 from theano import Type, Variable, Constant, tensor, config, scalar
 from theano.compile import SharedVariable
@@ -12,7 +16,7 @@ try:
     from pygpu import gpuarray
     from pygpu.elemwise import compare, elemwise2
 except ImportError:
-    pass
+    pygpu = None
 
 _context_reg = {}
 
@@ -74,7 +78,7 @@ def list_contexts():
 
 # Private method
 def _name_for_ctx(ctx):
-    for k, v in _context_reg:
+    for k, v in iteritems(_context_reg):
         if v == ctx:
             return k
         raise ContextNotDefined('context is not registered')
@@ -198,11 +202,14 @@ class GpuArrayType(Type):
                                   context=self.context)
         else:
             if not hasattr(data, 'dtype'):
-                # This is to convert objects that don't have a dtype
-                # (like lists).  We anticipate that the type below
-                # will match and we pass copy=False so it won't make a
-                # second object on the GPU.
-                data = gpuarray.array(data, copy=False, context=self.context)
+                converted_data = theano._asarray(data, self.dtype)
+                # We use the `values_eq` static function from TensorType
+                # to handle NaN values.
+                if TensorType.values_eq(numpy.asarray(data),
+                                        converted_data,
+                                        force_same_dtype=False):
+                    data = converted_data
+                    data = gpuarray.array(data, context=self.context)
 
             up_dtype = scalar.upcast(self.dtype, data.dtype)
             if up_dtype == self.dtype:
@@ -260,10 +267,10 @@ class GpuArrayType(Type):
         return GpuFromHost(self.context_name)(other)
 
     @staticmethod
-    def values_eq(a, b):
+    def values_eq(a, b, force_same_dtype=True):
         if a.shape != b.shape:
             return False
-        if a.typecode != b.typecode:
+        if force_same_dtype and a.typecode != b.typecode:
             return False
         a_eq_b = numpy.asarray(compare(a, '==', b))
         if a_eq_b.all():
@@ -751,3 +758,35 @@ Instance of :class:`GpuContextType` to use for the context_type
 declaration of an operation.
 """
 gpu_context_type = GpuContextType()
+
+
+# THIS WORKS But GpuArray instances don't compare equal to one
+# another, and what about __hash__ ?  So the unpickled version doesn't
+# equal the pickled version, and the cmodule cache is not happy with
+# the situation. The old back-end have this same comment and use the
+# same mechanism.
+def GpuArray_unpickler(npa, ctx_name):
+    if config.experimental.unpickle_gpu_on_cpu:
+        # directly return numpy array
+        warnings.warn(
+            "config.experimental.unpickle_gpu_on_cpu is set to True. "
+            "Unpickling GpuArray as numpy.ndarray")
+        return npa
+    elif pygpu:
+        ctx = get_context(ctx_name)
+        return pygpu.gpuarray.array(npa, copy=True, context=ctx)
+    else:
+        raise ImportError("pygpu not found. Cannot unpickle GpuArray")
+
+copyreg.constructor(GpuArray_unpickler)
+
+
+def GpuArray_pickler(cnda):
+    ctx_name = _name_for_ctx(cnda.context)
+    return (GpuArray_unpickler, (numpy.asarray(cnda), ctx_name))
+
+# In case pygpu is not imported.
+if pygpu is not None:
+    copyreg.pickle(pygpu.gpuarray.GpuArray,
+                   GpuArray_pickler,
+                   GpuArray_unpickler)