Convolutional LSTM - retain_graph Error

mayool: <p>Hey everyone,<br>
I am working on semantic segmentation and would like to extend an existing DeepLabV3+ (mobilenet backbone) with a recurrent unit (convolutional lstm).</p>
<p>My Idea was to concatinate the result of the segmentator at the current timestep T with its previous segmentation results (T-1 and T-2) and feed everything into the ConvLSTM (see picture).</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/0/e07c6cf0372afe94d231b1ba9eded97d67232177.png" data-download-href="https://discuss.pytorch.org/uploads/default/e07c6cf0372afe94d231b1ba9eded97d67232177" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/0/e07c6cf0372afe94d231b1ba9eded97d67232177_2_690x208.png" alt="image" data-base62-sha1="w1Tx6cvGwWTwDmydSvJXa0dvkOP" width="690" height="208" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/0/e07c6cf0372afe94d231b1ba9eded97d67232177_2_690x208.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/0/e07c6cf0372afe94d231b1ba9eded97d67232177_2_1035x312.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/0/e07c6cf0372afe94d231b1ba9eded97d67232177_2_1380x416.png 2x" data-dominant-color="E7E7E7"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1827×552 95.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>During training I always get this error:</p>
<blockquote>
<p>RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time</p>
</blockquote>
<p>Setting retain_graph to true however causes that cuda runs out of memory.</p>
<p>I’ve tried following <a href="https://discuss.pytorch.org/t/solved-training-a-simple-rnn/9055">this thread</a><br>
and it seems like the problem has to do with <code>hidden_state.detach()</code>.</p>
<p>My network looks like this:</p>
<pre><code class="lang-auto">class Deep_mobile_lstmV2(nn.Module):
    def __init__(self):
        super().__init__()
        self.base = deeplabv3plus_mobilenet(num_classes=2, pretrained_backbone=True)
        
        self.lstm = ConvLSTM(input_dim=2, hidden_dim=[2], kernel_size=(3, 3), num_layers=1, batch_first=True,
                             bias=True,
                             return_all_layers=False)
        self.hidden = None

    def forward(self, x, *args):
        # set old predictions 
        old_pred = args[0] # a list of the old prediction tensors

        # get semantic segmentation by
        out = self.base(x)
        out = out.unsqueeze(1) #[bs,channels,H,W] --&gt; [bs, timestep, channels, H,W]
        
        # set old predictions to 0 if they dont exists (for the first 2 frames)
        if len(args) != 0:
            if None in old_pred:
                for i in range(len(old_pred)):
                    old_pred[i] = torch.zeros_like(out)

            # concatinate old predictions with current predictions [_,1,_,_,_] --&gt; [_,3,_,_,_]
            out = [out] + old_pred
            out = torch.cat(out, dim =1)

        out, self.hidden = self.lstm(out, self.hidden)
        out = out[0][:,-1,:,:,:] # return only 1 segmentation
        return out
</code></pre>
<p>(Please note that I left out some unimportant code that would have asserted, that the shapes match for concatination, so it is easier to read)</p>
<p>The Deeplab implementation is from <a href="https://github.com/VainF/DeepLabV3Plus-Pytorch" rel="noopener nofollow ugc">this Git Repo</a> and the ConvLSTM is from <a href="https://github.com/ndrplz/ConvLSTM_pytorch" rel="noopener nofollow ugc">this repository</a> (with line 142 in convlstm.py replaced by hidden_state=hidden_state)</p>
<p>My training loop looks like this:</p>
<pre><code class="lang-auto">for epoch in range(num_epochs):
    old_pred = [None, None]
    for batch in train_loader:
        images, labels = batch
        pred = net(images, old_pred)
        loss = criterion(pred, labels.long())
        optimizer.zero_grad()
        # loss.backward(retain_graph=True) leads to memory issues
        loss.backward() # throws error
        optimizer.step()
        
        old_pred[1] = old_pred[0]
        old_pred[0] = pred.unsqueeze(1)
</code></pre>
<p>if I change<br>
<code>out, self.hidden = self.lstm(out, self.hidden)</code> into<br>
<code>out, self.hidden = self.lstm(out)</code> it is training but than the information in self.hidden would be lost with each new sample, right?</p>
<p>I am stuck at this problem for days and have no idea how to solve this. Any help would be appreciated!</p>

tetratrio: <p>So what happens when you dont detach the hidden state is that the previous computation graph is attached to it (if <code>retain_graph=True</code>). So when you calculate your next hidden state, the previous computation graph is added to the current computation graph. After a couple of iterations the complete graph is so big that you run out of memory.<br>
Not detaching will allow you to backprop through time (I will use BPTT as an abbreviation).</p>
<p>By instead detaching the hidden state you stop BPTT and the computation graph will not grow with each iteration. It is also possible to detach the hidden state at certain intervals to keep some BPTT.</p>
<pre><code class="lang-auto">out, self.hidden = self.lstm(out, self.hidden)
# If hidden is similar to LSTM it should be a tuple of two tensors
self.hidden = tuple(state.detach() for state in self.hidden)
</code></pre>
<p>This should run without needing to specify <code>retain_graph=True</code></p>

mayool: <p>Thanks for the answer! Yes, hidden is indeed a tuple of two tensors.</p>
<p>Unfortunately, this does not work and I really dont get why not…still getting the error from above:</p>
<aside class="quote no-group" data-username="mayool" data-post="1" data-topic="85200">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/41988e/48.png" class="avatar"> mayool:</div>
<blockquote>
<p>RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time</p>
</blockquote>
</aside>
<p>Does the position matter, where state.detach() is called?<br>
At the moment I call detach in the forward() method of my network.<br>
Does it need to be in the training loop?<br>
so i would call</p>
<pre><code class="lang-auto">new_hidden = None
for batch in train_loader:
        hidden_state = new_hidden if new_hidden is not None else init_hidden_with_zeros()
        images, labels = batch
        pred, new_hidden = net(images, hidden_state, old_pred)
        new_hidden = tuple(state.detach() for state in new_hidden)
        (...)
</code></pre>
<p>so basicially shifting the process of “taking care of the hidden state” from within the forward method into the train loop…</p>
<p>Does that change anything? <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>

tetratrio: <p>No problem, hopefully we can find the problem!</p>
<p>You just need to make sure you detach before using the hidden state again. It does not matter if you do it when you store the hidden state in the forward pass or when you retrieve the previous hidden state and feed it together with the new input to the LSTM.</p>
<p>The <code>old_pred</code> tensors need to be detached as well if it is concatenated with tensors in the current forward pass. Because the <code>old_pred</code> tensors still have a computation graph!</p>
<pre><code class="lang-auto"># In training loop
old_pred[1] = old_pred[0]
old_pred[0] = pred.unsqueeze(1).detach()
</code></pre>
<p>Hopefully it works this way. Gonna be unavailable soon so might take a day for me to respond next time <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>

mayool: <p>Mate you saved my day! detaching the pred works! Thank you so much <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>
<p>Did I get this right that know backpropagation is only happening for the new Timestep T and not for T-1, T-2?</p>

tetratrio: <p>Awesome, glad I could help!</p>
<p>You have two recurrent parts, the previous prediction as well as the hidden state of the LSTM. Detaching these will stop gradients flowing back in time, so you only get gradients for the current timestep.</p>
<p>If you want to backprop in time you can always detach at a certain interval! For this to work correctly I think you need to do something like this:</p>
<pre><code class="lang-auto">detach_interval = 5
for i, batch in enumerate(data):
    input, target = batch
    pred = model(input)
    loss = criterion(pred, target)
    if i &gt; 0 and i % detach_interval == 0:
        loss.backward(retain_graph=False)
        # Here you need to detach tensors that will be reused next forward pass
        ...
    else:
        loss.backward(retain_graph=True)
    optimizer.step()
</code></pre>
<p>Note that you need to retain the computation graph when calling <code>backward(retain_graph=True)</code> if you are going to reuse tensors for next forward pass without detaching them. On the last <code>backward()</code> before detaching, you can let the backprop clean out the graph.</p>
<p>Hopefully this makes sense!</p>

mayool: <p>Makes sense to me,<br>
thank you once more!<br>
Recurrency is still confusing for me <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>
<p>Hope you have a great weekend!</p>