Implementing Truncated Backpropagation Through Time

adrianjav: <p>Hello,<br>
I’m implementing a recursive network that is going to be trained with very long sequences. I had memory problems when training because of that excessive length and I decided to use a <strong>truncated</strong>-BPTT algorithm to train it as described <a href="https://machinelearningmastery.com/gentle-introduction-backpropagation-time/" rel="nofollow noopener">here</a>, that is,</p>
<p><code>every k1 steps backpropagate taking k2 back steps</code></p>
<p>checking some <a href="https://github.com/pytorch/examples/blob/master/word_language_model/main.py" rel="nofollow noopener">examples</a>  I could easily write the case when <code>k1 = k2</code>. However, I haven’t been able to implement the general case yet.</p>
<p>My first idea was to freeze the  gradient graph after the first <code>k2</code> steps and keep changing the <code>Variables</code> referenced there. Later, I saw that the graph doesn’t directly reference earlier Variables, instead it directly incorporates the gradient graphs from that <code>Variables</code>. So I thought about looking for that copied subgraphs and substituting them by a <code>Variable</code> reference, but I found out that the gradient graph is <em>unwritable</em> and <em>immutable</em>.</p>
<p>The only idea I have left is to recompute the calculations in the overlaps between backpropagations. That’s going to work, but I’d really love to avoid recomputing things.</p>
<p>Any idea on how to implement this efficiently?</p>

albanD: <p>Just to be sure before I write something long that is not what you asked:</p>
<ul>
<li>You have an <code>nn.Module</code> (lets call it <code>one_step_module</code>) that does one step given a current state, an input and produce an output</li>
<li>You have another <code>nn.Module</code> (lets call it <code>loss_module</code>) that given the output and a target gives you the loss for this output.</li>
</ul>
<p>And in pseudo code, what you want is:</p>
<pre><code class="lang-auto">state = init_state
for i, (inp, target) in enumerate(my_very_long_sequence_of_inputs):
    output, state = one_step_module(inp, state)
    if (i+1)%k1 == 0:
        loss = loss_module(output, target)
        # You want the function below
        loss.backward_only_k2_last_calls_to_one_step_module()
</code></pre>

adrianjav: <p>That’s pretty much what I want to do, indeed. So far I just detach the variables inside that if statement, so I achieve the effect of having <code>k1 = k2</code>.</p>
<p>Maybe I should have put some pseudocode from the very beginning. Sorry about that.</p>

albanD: <p>Here is an implementation that will work for any k1 and k2 and will reduce memory usage as much as possible.<br>
If k2 is not huge and the <code>one_step_module</code> is relatively big, the slowdown of doing multiple backward should be negligible.</p>
<p>The code is not super clean and has been tested only against current master branch (where Variable and Tensor are merged) so you might need slight modifications if you use 0.3.<br>
Hope this helps.</p>
<pre><code class="lang-auto">
class TBPTT():
    def __init__(self, one_step_module, loss_module, k1, k2, optimizer):
        self.one_step_module = one_step_module
        self.loss_module = loss_module
        self.k1 = k1
        self.k2 = k2
        self.retain_graph = k1 &lt; k2
        # You can also remove all the optimizer code here, and the
        # train function will just accumulate all the gradients in
        # one_step_module parameters
        self.optimizer = optimizer

    def train(self, input_sequence, init_state):
        states = [(None, init_state)]
        for j, (inp, target) in enumerate(input_sequence):

            state = states[-1][1].detach()
            state.requires_grad=True
            output, new_state = self.one_step_module(inp, state)
            states.append((state, new_state))

            while len(states) &gt; self.k2:
                # Delete stuff that is too old
                del states[0]

            if (j+1)%self.k1 == 0:
                loss = self.loss_module(output, target)

                optimizer.zero_grad()
                # backprop last module (keep graph only if they ever overlap)
                start = time.time()
                loss.backward(retain_graph=self.retain_graph)
                for i in range(self.k2-1):
                    # if we get all the way back to the "init_state", stop
                    if states[-i-2][0] is None:
                        break
                    curr_grad = states[-i-1][0].grad
                    states[-i-2][1].backward(curr_grad, retain_graph=self.retain_graph)
                print("bw: {}".format(time.time()-start))
                optimizer.step()



seq_len = 20
layer_size = 50

idx = 0

class MyMod(nn.Module):
    def __init__(self):
        super(MyMod, self).__init__()
        self.lin = nn.Linear(2*layer_size, 2*layer_size)

    def forward(self, inp, state):
        global idx
        full_out = self.lin(torch.cat([inp, state], 1))
        # out, new_state = full_out.chunk(2, dim=1)
        out = full_out.narrow(1, 0, layer_size)
        new_state = full_out.narrow(1, layer_size, layer_size)
        def get_pr(idx_val):
            def pr(*args):
                print("doing backward {}".format(idx_val))
            return pr
        new_state.register_hook(get_pr(idx))
        out.register_hook(get_pr(idx))
        print("doing fw {}".format(idx))
        idx += 1
        return out, new_state


one_step_module = MyMod()
loss_module = nn.MSELoss()
input_sequence = [(torch.rand(200, layer_size), torch.rand(200, layer_size))] * seq_len

optimizer = torch.optim.SGD(one_step_module.parameters(), lr=1e-3)

runner = TBPTT(one_step_module, loss_module, 5, 7, optimizer)

runner.train(input_sequence, torch.zeros(200, layer_size))
print("done")
</code></pre>

adrianjav: <p>So the idea of your code is to isolate the variables in each time-step and every <code>k1</code> steps <em>“rewire”</em> the last <code>k2</code> states, right? I like it, I should’ve found it out myself to be fair.</p>
<p>And I agree that as long as <code>k2 - k1</code> isn’t too big the overhead should be negligible.</p>
<p>Thanks for your help <a class="mention" href="/u/alband">@albanD</a>! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

albanD: <p>Yes, that is exactly the idea !</p>
<p>Happy it helps <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

riccardosamperna: <p>I have a very similar problem where I am trying to unroll a recurrent neural network, in my case I don’t use truncated backprop but just BPTT. The network takes as input the previous output and my training code looks like:</p>
<pre><code class="lang-auto">for epoch in range(1, args.epochs + 1):
        model.train()
        epoch_loss = 0
        for i, sequence in enumerate(training_data_loader):
            optimizer.zero_grad()
            loss = 0

            output = torch.zeros(sequence['input'].size(0), 3, sequence['input'].size(3), sequence['input'].size(4)).cuda(args.gpu, non_blocking=True)

            for j in range(sequence['input'].size(2) - 1):
                inputs = torch.index_select(sequence['input'], 2, torch.tensor([j,j+1])).cuda(args.gpu, non_blocking=True)
                t = torch.squeeze(torch.index_select(sequence['target'], 2, torch.tensor([j+1])), 2).cuda(args.gpu, non_blocking=True)
                output, l = model(inputs, output, t, i, writer, im_out)
                loss += l 
            
            loss.backward()
            optimizer.step()
</code></pre>
<p>It looks like that the gradient is not flowing backwards, do you know what the issue could be?</p>

vfdev-5: <p>Just for info, in Ignite today we also provide a <a href="https://github.com/pytorch/ignite/blob/master/ignite/contrib/engines.py#L43" rel="nofollow noopener">trainer implementation of TBPTT</a></p>

adrianjav: <p>That’s a neat implementation! But that’s a particular case of TBPTT since it assumes that <code>k1 = k2</code> (tbtt_step in this case).</p>
<p>My implementation had two parameters, k1 and k2, and I found out that there’s a nice improvement when <code>k2 &gt; k1</code>. I’d love to contribute when I had some time.</p>

vfdev-5: <blockquote>
<p>My implementation had two parameters, k1 and k2, and I found out that there’s a nice improvement when <code>k2 &gt; k1</code> . I’d love to contribute when I had some time.</p>
</blockquote>
<p><a class="mention" href="/u/adrianjav">@adrianjav</a> thanks! PRs are very welcome <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

gbarlas: <p>For k1=3 and k2=5 this code back-propagates as shown below:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/2X/0/0e838103463bb1b28bbd8123d81e1f9d8756c293.gif" alt="wrong" data-base62-sha1="24oqGTjPEa11fkIMWP6wcf3TBqr" width="530" height="112" class="animated"></p>

gbarlas: <p>(continue …)<br>
But it should back-propagate like this:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/2X/e/e805d58fc9b46ec07aa09f00c435bf20ca332c54.gif" alt="correct" data-base62-sha1="x6zgV55OxpSIiQ3976KnLcgEnqI" width="530" height="112" class="animated"></p>
<p>It’s easy to implement the second one by changing few lines in <em>train</em> function.</p>
<pre><code class="lang-auto">  def train(self, input_sequence, init_state):
        states = [(None, init_state)]

        outputs = []
        targets = []

        for i, (inp, target) in enumerate(input_sequence):

            state = states[-1][1].detach()
            state.requires_grad=True
            output, new_state = self.one_step_module(inp, state)

            outputs.append(output)
            targets.append(target)
            while len(outputs) &gt; self.k1:
                # Delete stuff that is too old
                del outputs[0]
                del targets[0]

            states.append((state, new_state))
            while len(states) &gt; self.k2:
                # Delete stuff that is too old
                del states[0]
            if (i+1)%self.k1 == 0:
                # loss = self.loss_module(output, target)

                optimizer.zero_grad()
                # backprop last module (keep graph only if they ever overlap)
                start = time.time()
                # loss.backward(retain_graph=self.retain_graph)
                for j in range(self.k2-1):

                    if j &lt; self.k1:
                        loss = self.loss_module(outputs[-j-1], targets[-j-1])
                        loss.backward(retain_graph=True)

                    # if we get all the way back to the "init_state", stop
                    if states[-j-2][0] is None:
                        break
                    curr_grad = states[-j-1][0].grad
                    states[-j-2][1].backward(curr_grad, retain_graph=self.retain_graph)
                print("bw: {}".format(time.time()-start))
                optimizer.step()
</code></pre>
<p>I think it may also be possible to implement it without <em>outputs</em> and <em>targets</em> lists.</p>
<p>I splitted my post, because as a new user i can post only one image per post.</p>

adrianjav: <p>I think that’s not how is suppose to work either. As I understood the algorithm, <code>k1</code> only refers to the steps the model does before actually doing the backpropagation. <code>k2</code> refers to the length of that backpropagation.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/2X/d/df14ddfb70762f9669341a3c088397701c8efe4d.png" alt="Screenshot%20from%202018-09-13%2018-14-35" data-base62-sha1="vPtbno9Ny1lMALyTVSMFhdxs5g9" width="565" height="192"></p>
<p>I might be mistaken, but you can always look at better sources. The code shown is from <a href="https://www.cs.utoronto.ca/~ilya/pubs/ilya_sutskever_phd_thesis.pdf" rel="noopener nofollow ugc">this phd thesis</a> (section 2.8.6).</p>

satyakrishnagorti: <p>Hey, just to confirm, I think you have two loops with the same index variable i. This could potentially be wrong, unless it was intended this way, which I don’t think should be the case.</p>

albanD: <p>Yes you are right this is not good.<br>
In this particular case it does not change anything but could have.<br>
I modified the original post to use j for one of the two loops.</p>

DozyC: <aside class="quote no-group" data-username="vfdev-5" data-post="8" data-topic="15500">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/vfdev-5/48/3535_2.png" class="avatar"> vfdev-5:</div>
<blockquote>
<p>Just for info, in Ignite today we also provide a <a href="https://github.com/pytorch/ignite/blob/master/ignite/contrib/engines.py#L43" rel="noopener nofollow ugc">trainer implementation of TBPTT</a></p>
</blockquote>
</aside>
<p>Here is an updated link</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/ignite/blob/master/ignite/contrib/engines/tbptt.py">
  <header class="source">

      <a href="https://github.com/pytorch/ignite/blob/master/ignite/contrib/engines/tbptt.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/ignite/blob/master/ignite/contrib/engines/tbptt.py" target="_blank" rel="noopener nofollow ugc">pytorch/ignite/blob/master/ignite/contrib/engines/tbptt.py</a></h4>


      <pre><code class="lang-py"># coding: utf-8
import collections.abc as collections
from typing import Callable, Mapping, Optional, Sequence, Union

import torch
import torch.nn as nn
from torch.optim.optimizer import Optimizer

from ignite.engine import _prepare_batch, Engine, EventEnum
from ignite.utils import apply_to_tensor


class Tbptt_Events(EventEnum):
    """Aditional tbptt events.

    Additional events for truncated backpropagation throught time dedicated
    trainer.
    """

    TIME_ITERATION_STARTED = "time_iteration_started"
</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/ignite/blob/master/ignite/contrib/engines/tbptt.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


Ceyer_Wakilpoor: <aside class="quote no-group" data-username="albanD" data-post="4" data-topic="15500">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>states.append((state, new_state))</p>
</blockquote>
</aside>
<p>this might be a stupid question, but did you need to store tuples, couldn’t you have just stored a list of old states?</p>

albanD: <p>Hi,</p>
<p>Yes you can, but since all of these are just references to the actual Tensors, it’s not a problem to store each of them multiple times.<br>
So whichever your more comfortable with !</p>

Ceyer_Wakilpoor: <p>this is a way to do back prop, but the original is not incorrect, it depends on the problem. If you have a label at the end of a sequence then your solution is not applicable, but if you are update the parameters for a decoder, this is useful, but one potential issue (that i’m currently trying to figure out) is that won’t your hidden state only propagate the error from the last label, since the states are being detached? The gradient accumulation will be on the last hidden state therefore when you perform back prop on previous labels those will only be factored into the label’s respective timestep. Or so it seems to me, or will the older state still get gradient accumulation from previous label errors ? I might be assuming the wrong behavior about detach()</p>

gbarlas: <p>This post [<a href="https://discuss.pytorch.org/t/truncated-backprop-data-clarification/34854/2" class="inline-onebox">Truncated backprop data clarification</a>] may help you.</p>