Repository: Theano/Theano

Issue:
Generalize Scan memory/speed trade-off
https://github.com/Theano/Theano/issues/4124 issue was about adding that.

Current implementation have this restriction that we want to lift:
- [ ] The value of `save_every_N` divides the number of steps the scan will run without remainder.

I'm not sure we should lift this limit. What do you think?
- Only singly-recurrent and non-recurrent outputs are used. No multiple recurrences.



Discussion:
The first thing to try to implement is to pass just the slice to the inner scan with the right shape. If all goes well, it will work. But maybe something will request that the inner scan have the same number of step. In that case, we could pad in the outer scan the slice to always have the same size.

@nouiz Just a quick note that we discussed with @lamblin and we have no choice but padding the input because of the reshape of the sequences:

``` python
o_sequences = [s.reshape([s.shape[0] / save_every_N, save_every_N] +
                           [s.shape[i] for i in range(1, s.ndim)],
                          s.ndim + 1) for s in sequences]
```

Here we definitively need to have sequences with length that are a multiple of `save_every_N`. So we will pad the input sequence.

ok. I think we should pad by default, but add a parameter to not pad, so
that people that can garanty the shapes are good, won't pay this price.

Should we do such a pad op? This appear at a few different places and
having an op in the graph that would garanty we have it optimized (return a
view if no pad needed) would be very useful. I don't think we can rely on
ifelse to do that well in all cases without side effect that aren't fun.

On Thu, Nov 17, 2016 at 5:31 PM, César Laurent notifications@github.com
wrote:

> @nouiz https://github.com/nouiz Just a quick note that we discussed
> with @lamblin https://github.com/lamblin and we have no choice but
> padding the input because of the reshape of the sequences:
> 
> o_sequences = [s.reshape([s.shape[0] / save_every_N, save_every_N] +
>                            [s.shape[i] for i in range(1, s.ndim)],
>                           s.ndim + 1) for s in sequences]
> 
> Here we definitively need to have sequences with length that are a
> multiple of save_every_N. So we will pad the input sequence.
> 
> —
> You are receiving this because you were mentioned.
> Reply to this email directly, view it on GitHub
> https://github.com/Theano/Theano/issues/5106#issuecomment-261390824, or mute
> the thread
> https://github.com/notifications/unsubscribe-auth/AALC-5IUp3R71NNi8PzESED7XIlpna4nks5q_NW9gaJpZM4KYwbd
> .

I'm thinking about that again. SetSubtensor and variant, if they see that we override the full tensor, could return a view on the original data (if there is no broadcast too).

This would request that we do special code to force those nodes to be inplace on the value instead of the original data or change Theano to allow one output to be a view of many inputs.

Between those 3 options, the new op or added a new parameter to setsubtensor to be inplace on the value would be faster to implement then the last one.

What about @Thrandis implementing the not optimized code for now and make a new issue about adding that is needed for the optimized version in Theano?

@nouiz I'm ok with doing anything. It is true that since we do need padding, we should have an op for it!

I think we should start with the un-optimized code so that it works and the interface is stable.
Then we can see how to optimize that memory-wise.

When we decide to optimize, I like the idea of SetSubtensor being inplace on the value. Another option may be to optimize the runtime of Join (instead of SetSubtensor) so that it can work inplace on an input if it is the only one with a nonzero size. All of these require to be able to work inplace on different inputs, though, which is its own can of worms.

The SetSubtensor\* won't work. It will only save the copy, it won't save the creation of a big tensor of zeros.

@Thrandis tried to use join. If we allow join to concatenate empty tensor (it crash now in that cases) then that would work. We could do the quick fix to make it inplace on the first inputs. So it don't ask the need to be inplace on many inputs.

This would reuse existing ops instead of creating a new ops. But this won't be a different interface then to numpy as we think numpy don't allow to do that, as we call numpy.concatenate.


Resolving patch:
diff --git a/theano/scan_module/scan.py b/theano/scan_module/scan.py
--- a/theano/scan_module/scan.py
+++ b/theano/scan_module/scan.py
@@ -265,11 +265,9 @@ def f(x):
         ``n_steps`` is the number of steps to iterate given as an int
         or Theano scalar. If any of the input sequences do not have
         enough elements, scan will raise an error. If the *value is 0* the
-        outputs will have *0 rows*. If the value is negative, ``scan``
-        will run backwards in time. If the ``go_backwards`` flag is already
-        set and also ``n_steps`` is negative, ``scan`` will run forward
-        in time. If n_steps is not provided, ``scan`` will figure
-        out the amount of steps it should run given its input sequences.
+        outputs will have *0 rows*. If n_steps is not provided, ``scan`` will
+        figure out the amount of steps it should run given its input
+        sequences. ``n_steps`` < 0 is not supported anymore.
 
     truncate_gradient
         ``truncate_gradient`` is the number of steps to use in truncated
diff --git a/theano/scan_module/scan_checkpoints.py b/theano/scan_module/scan_checkpoints.py
--- a/theano/scan_module/scan_checkpoints.py
+++ b/theano/scan_module/scan_checkpoints.py
@@ -4,7 +4,8 @@
 
 
 def scan_checkpoints(fn, sequences=[], outputs_info=None, non_sequences=[],
-                     name="checkpointscan_fn", n_steps=None, save_every_N=10):
+                     name="checkpointscan_fn", n_steps=None, save_every_N=10,
+                     padding=True):
     """Scan function that uses less memory, but is more restrictive.
 
     In :func:`~theano.scan`, if you compute the gradient of the output
@@ -52,19 +53,23 @@ def scan_checkpoints(fn, sequences=[], outputs_info=None, non_sequences=[],
 
     n_steps
         ``n_steps`` is the number of steps to iterate given as an int
-        or Theano scalar. If any of the input sequences do not have
-        enough elements, scan will raise an error. If the **value is 0**
-        the outputs will have **0 rows**. If the value is negative,
-        ``scan`` will run backwards in time. If the ``go_backwards`` flag
-        is already set and also ``n_steps`` is negative, ``scan`` will run
-        forward in time. If n_steps is not provided, ``scan`` will figure
-        out the amount of steps it should run given its input sequences.
+        or Theano scalar (> 0). If any of the input sequences do not have
+        enough elements, scan will raise an error. If n_steps is not provided,
+        ``scan`` will figure out the amount of steps it should run given its
+        input sequences.
 
     save_every_N
         ``save_every_N`` is the number of steps to go without storing
         the computations of ``scan`` (ie they will have to be recomputed
         during the gradient computation).
 
+    padding
+        If the length of the sequences is not a multiple of ``save_every_N``,
+        the sequences will be zero padded to make this version of ``scan``
+        work properly, but will also result in a memory copy. It can be
+        avoided by setting ``padding`` to False, but you need to make
+        sure the length of the sequences is a multple of ``save_every_N``.
+
     Returns
     -------
     tuple
@@ -96,16 +101,31 @@ def scan_checkpoints(fn, sequences=[], outputs_info=None, non_sequences=[],
     if n_steps is None:
         n_steps = sequences[0].shape[0]
 
-    # Compute the number of steps of the inner and of the outer scan
-    o_n_steps = theano.tensor.cast(n_steps / save_every_N, 'int64')
-    i_n_steps = save_every_N
+    # Compute the number of steps of the outer scan
+    o_n_steps = theano.tensor.cast(theano.tensor.ceil(n_steps / save_every_N),
+                                   'int64')
+
+    # Compute the number of steps of the inner scan
+    i_n_steps = save_every_N * theano.tensor.ones((o_n_steps,), 'int64')
+    mod = n_steps % save_every_N
+    last_n_steps = theano.tensor.switch(theano.tensor.eq(mod, 0),
+                                        save_every_N, mod)
+    i_n_steps = theano.tensor.set_subtensor(i_n_steps[-1], last_n_steps)
+
+    # Pad the sequences if needed
+    if padding:
+        for i, s in enumerate(sequences):
+            n = s.shape[0] % save_every_N
+            z = theano.tensor.zeros((n, s.shape[1:]), dtype=s.dtype)
+            sequences[i] = theano.tensor.concatenate([s, z], axis=0)
 
     # Establish the input variables of the outer scan
     o_sequences = [s.reshape([s.shape[0] / save_every_N, save_every_N] +
                              [s.shape[i] for i in range(1, s.ndim)],
                              s.ndim + 1) for s in sequences]
+    o_sequences.append(i_n_steps)
     new_nitsots = [i for i in outputs_info if i is None]
-    o_nonsequences = non_sequences + [i_n_steps]
+    o_nonsequences = non_sequences
 
     def outer_step(*args):
         # Separate the received arguments into their respective (seq, outputs
@@ -117,11 +137,11 @@ def outer_step(*args):
 
         # Call the user-provided function with the proper arguments
         results, updates = theano.scan(fn=fn,
-                                       sequences=i_sequences,
+                                       sequences=i_sequences[:-1],
                                        outputs_info=i_outputs_infos,
-                                       non_sequences=i_non_sequences[:-1],
+                                       non_sequences=i_non_sequences,
                                        name=name + "_inner",
-                                       n_steps=i_non_sequences[-1])
+                                       n_steps=i_sequences[-1])
         if not isinstance(results, list):
             results = [results]
 


Test patch:
diff --git a/theano/scan_module/tests/test_scan_checkpoints.py b/theano/scan_module/tests/test_scan_checkpoints.py
--- a/theano/scan_module/tests/test_scan_checkpoints.py
+++ b/theano/scan_module/tests/test_scan_checkpoints.py
@@ -38,14 +38,14 @@ def test_forward_pass(self):
         """Test forward computation of A**k."""
         f = theano.function(inputs=[self.A, self.k],
                             outputs=[self.result, self.result_check])
-        out, out_check = f(range(10), 100)
+        out, out_check = f(range(10), 101)
         assert numpy.allclose(out, out_check)
 
     def test_backward_pass(self):
         """Test gradient computation of A**k."""
         f = theano.function(inputs=[self.A, self.k],
                             outputs=[self.grad_A, self.grad_A_check])
-        out, out_check = f(range(10), 100)
+        out, out_check = f(range(10), 101)
         assert numpy.allclose(out, out_check)
 
     @unittest.skipUnless(PYGPU_AVAILABLE, 'Requires pygpu.')