Merge pull request #2998 from nouiz/seed_mrg_rng

theano/sandbox/rng_mrg.py

@@ -1077,28 +1077,68 @@ class MRG_RandomStreams(object):
         self.state_updates = []
 
         super(MRG_RandomStreams, self).__init__()
+
+        # Needed to reset the streams.
+        self.default_instance_seed = seed
+
+        self.set_rstate(seed)
+
+        if use_cuda is None:
+            self.use_cuda = cuda_enabled
+        else:
+            self.use_cuda = use_cuda
+
+    def set_rstate(self, seed):
         if isinstance(seed, int):
             if seed == 0:
                 raise ValueError('seed should not be 0', seed)
             elif seed >= M2:
                 raise ValueError('seed should be less than %i' % M2, seed)
-            self.rstate = numpy.asarray([seed]*6, dtype='int32')
+            self.rstate = numpy.asarray([seed] * 6, dtype='int32')
         elif len(seed) == 6:
             if seed[0] == 0 and seed[1] == 0 and seed[2] == 0:
-                raise ValueError('The first 3 values of seed should not be all 0', seed)
+                raise ValueError(
+                    'The first 3 values of seed should not be all 0', seed)
             if seed[3] == 0 and seed[4] == 0 and seed[5] == 0:
-                raise ValueError('The last 3 values of seed should not be all 0', seed)
+                raise ValueError(
+                    'The last 3 values of seed should not be all 0', seed)
             if seed[0] >= M1 or seed[1] >= M1 or seed[2] >= M1:
-                raise ValueError('The first 3 values of seed should be less than %i' % M1, seed)
+                raise ValueError(
+                    'The first 3 values of seed should be less than %i' % M1,
+                    seed)
             if seed[3] >= M2 or seed[4] >= M2 or seed[5] >= M2:
-                raise ValueError('The last 3 values of seed should be less than %i' % M2, seed)
+                raise ValueError(
+                    'The last 3 values of seed should be less than %i' % M2,
+                    seed)
             self.rstate = numpy.asarray(seed, dtype='int32')
         else:
             raise TypeError("seed should be 1 integer or 6 integers")
-        if use_cuda is None:
-            self.use_cuda = cuda_enabled
-        else:
-            self.use_cuda = use_cuda
+
+    def seed(self, seed=None):
+        """Re-initialize each random stream
+
+        :param seed: each random stream will be assigned a unique
+        state that depends deterministically on this value.
+
+        :type seed: None or integer in range 0 to 2**30
+
+        :rtype: None
+
+        """
+        if seed is None:
+            seed = self.default_instance_seed
+        self.set_rstate(seed)
+
+        for old_r, new_r, size, nstreams in self.state_updates:
+            if nstreams is None:
+                nstreams = self.n_streams(size)
+            rstates = self.get_substream_rstates(nstreams,
+                                                 new_r.owner.outputs[1].dtype)
+            assert (old_r.get_value(borrow=True,
+                                    return_internal_type=True).shape ==
+                    rstates.shape)
+            assert rstates.dtype == old_r.dtype
+            old_r.set_value(rstates, borrow=True)
 
     def inc_rstate(self):
         """Update self.rstate to be skipped 2^134 steps forward to the next stream start"""
@@ -1106,10 +1146,11 @@ class MRG_RandomStreams(object):
         self.rstate = multMatVect(self.rstate, A1p134, M1, A2p134, M2)
         assert self.rstate.dtype == numpy.int32
 
-    def get_substream_rstates(self, n_streams, inc_rstate=True):
+    def get_substream_rstates(self, n_streams, dtype, inc_rstate=True):
         """Initialize a matrix in which each row is a MRG stream state,
         and they are spaced by 2**72 samples.
         """
+        assert isinstance(dtype, str)
         assert n_streams < 2**72
         assert n_streams > 0
         rval = numpy.zeros((n_streams, 6), dtype='int32')
@@ -1136,15 +1177,25 @@ class MRG_RandomStreams(object):
 
         if inc_rstate:
             self.inc_rstate()
+        if self.use_cuda and dtype == 'float32':
+            rval = rval.flatten()
+            # HACK - we use fact that int32 and float32 have same size to
+            # sneak ints into the CudaNdarray type.
+            # these *SHOULD NEVER BE USED AS FLOATS*
+            tmp_float_buf = numpy.frombuffer(rval.data, dtype='float32')
+            assert tmp_float_buf.shape == rval.shape
+            assert (tmp_float_buf.view('int32') == rval).all()
+            rval = tmp_float_buf
+
         return rval
 
     def n_streams(self, size):
         return guess_n_streams(size)
 
-    def pretty_return(self, node_rstate, new_rstate, sample):
+    def pretty_return(self, node_rstate, new_rstate, sample, size, nstreams):
         sample.rstate = node_rstate
         sample.update = (node_rstate, new_rstate)
-        self.state_updates.append((node_rstate, new_rstate))
+        self.state_updates.append((node_rstate, new_rstate, size, nstreams))
         node_rstate.default_update = new_rstate
         return sample
 
@@ -1201,21 +1252,13 @@ class MRG_RandomStreams(object):
                 raise TypeError("size must be a tuple of int or a Theano "
                                 "Variable with 1 dimension, got " + str(size) +
                                 " of type " + str(type(size)))
-
+        orig_nstreams = nstreams
         if nstreams is None:
             nstreams = self.n_streams(size)
+        rstates = self.get_substream_rstates(nstreams, dtype)
 
         if self.use_cuda and dtype == 'float32':
-            rstates = self.get_substream_rstates(nstreams)
-            rstates = rstates.flatten()
-            # HACK - we use fact that int32 and float32 have same size to
-            # sneak ints into the CudaNdarray type.
-            # these *SHOULD NEVER BE USED AS FLOATS*
-            tmp_float_buf = numpy.frombuffer(rstates.data, dtype='float32')
-            assert tmp_float_buf.shape == rstates.shape
-            assert (tmp_float_buf.view('int32') == rstates).all()
-            # transfer to device
-            node_rstate = float32_shared_constructor(tmp_float_buf)
+            node_rstate = float32_shared_constructor(rstates)
             assert isinstance(node_rstate.type, CudaNdarrayType)
 
             # we can't use the normal mrg_uniform constructor + later
@@ -1225,12 +1268,14 @@ class MRG_RandomStreams(object):
             # reinterpretation.
             u = self.pretty_return(node_rstate,
                                    *GPU_mrg_uniform.new(node_rstate,
-                                                        ndim, dtype, size))
+                                                        ndim, dtype, size),
+                                   size=size, nstreams=orig_nstreams)
         else:
-            node_rstate = shared(self.get_substream_rstates(nstreams))
+            node_rstate = shared(rstates)
             u = self.pretty_return(node_rstate,
                                    *mrg_uniform.new(node_rstate,
-                                                    ndim, dtype, size))
+                                                    ndim, dtype, size),
+                                   size=size, nstreams=orig_nstreams)
         # Add a reference to distinguish from other shared variables
         node_rstate.tag.is_rng = True
         r = u * (high - low) + low

theano/sandbox/test_rng_mrg.py

@@ -932,6 +932,57 @@ def test_multMatVect():
     assert numpy.allclose(r_a2, r_b[3:])
 
 
+def test_seed_fn():
+    test_use_cuda = [False]
+    if cuda_available:
+        test_use_cuda.append(True)
+    idx = tensor.ivector()
+    for use_cuda in test_use_cuda:
+        if config.mode == 'FAST_COMPILE' and use_cuda:
+            mode = 'FAST_RUN'
+        else:
+            mode = config.mode
+
+        for new_seed, same in [(234, True), (None, True), (23, False)]:
+            random = MRG_RandomStreams(234, use_cuda=use_cuda)
+            fn1 = theano.function([], random.uniform((2, 2), dtype='float32'),
+                                  mode=mode)
+            fn2 = theano.function([], random.uniform((3, 3), nstreams=2,
+                                                     dtype='float32'),
+                                  mode=mode)
+            fn3 = theano.function([idx],
+                                  random.uniform(idx, nstreams=3, ndim=1,
+                                                 dtype='float32'),
+                                  mode=mode)
+
+            fn1_val0 = fn1()
+            fn1_val1 = fn1()
+            assert not numpy.allclose(fn1_val0, fn1_val1)
+            fn2_val0 = fn2()
+            fn2_val1 = fn2()
+            assert not numpy.allclose(fn2_val0, fn2_val1)
+            fn3_val0 = fn3([4])
+            fn3_val1 = fn3([4])
+            assert not numpy.allclose(fn3_val0, fn3_val1)
+            assert fn1_val0.size == 4
+            assert fn2_val0.size == 9
+
+            random.seed(new_seed)
+
+            fn1_val2 = fn1()
+            fn1_val3 = fn1()
+            fn2_val2 = fn2()
+            fn2_val3 = fn2()
+            fn3_val2 = fn3([4])
+            fn3_val3 = fn3([4])
+            assert numpy.allclose(fn1_val0, fn1_val2) == same
+            assert numpy.allclose(fn1_val1, fn1_val3) == same
+            assert numpy.allclose(fn2_val0, fn2_val2) == same
+            assert numpy.allclose(fn2_val1, fn2_val3) == same
+            assert numpy.allclose(fn3_val0, fn3_val2) == same
+            assert numpy.allclose(fn3_val1, fn3_val3) == same
+
+
 if __name__ == "__main__":
     rng = MRG_RandomStreams(numpy.random.randint(2147462579))
     import time