Backward() to compute partial derivatives without retain_graph= True

Stefano_Savian: <p>Dear all, i am implementing the paper GradNorm, I have found 2 implementations but both of them require a lot of memory.</p>
<p>Let’s suppose i have a loss L and a network W:<br>
L = loss1 + loss2 + lossN<br>
W are the weights with grad = True</p>
<p>I want to do :<br>
(dloss1)/dW<br>
save gradient norm<br>
(dloss2)/dW<br>
save gradient norm<br>
(dL)/dW<br>
update parameters</p>
<p>An example of  implementation can be found at:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/falkaer/artist-group-factors/blob/master/train_mtn.py">
  <header class="source">

      <a href="https://github.com/falkaer/artist-group-factors/blob/master/train_mtn.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/falkaer/artist-group-factors/blob/master/train_mtn.py" target="_blank" rel="noopener nofollow ugc">falkaer/artist-group-factors/blob/master/train_mtn.py</a></h4>


      <pre><code class="lang-py">import torch
import torch.nn as nn
import torch.nn.functional as F

from features import get_data_loaders, genre_counts, FramedFeatureDataset
from model import MTN, MTNFC
from logger import ModelLogger

class SoftCrossEntropyLoss(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, inputs, soft_targets):
        log_likelihood = F.log_softmax(inputs, dim=1)
        cross_entropy = -torch.sum(soft_targets * log_likelihood, dim=1)
        
        return torch.mean(cross_entropy)  # mean across batch

class MTNFCTrainer:
    
</code></pre>



  This file has been truncated. <a href="https://github.com/falkaer/artist-group-factors/blob/master/train_mtn.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

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

  <div style="clear: both"></div>
</aside>
<p>
the interesting lines are:</p>
<pre><code class="lang-auto">    # compute and retain gradients
    total_weighted_loss.backward(retain_graph=True)
    
    # GRADNORM - learn the weights for each tasks gradients
    
    # zero the w_i(t) gradients since we want to update the weights using gradnorm loss
    self.weights.grad = 0.0 * self.weights.grad
    
    W = list(self.model.mtn.shared_block.parameters())
    norms = []
    
    for w_i, L_i in zip(self.weights, task_losses):
        # gradient of L_i(t) w.r.t. W
        gLgW = torch.autograd.grad(L_i, W, retain_graph=True)
        
        # G^{(i)}_W(t)
        norms.append(torch.norm(w_i * gLgW[0]))
    
    norms = torch.stack(norms)
</code></pre>
<p>BUT I want to avoid retain_graph = True as this leads to out of memory for my network.</p>
<p>A similar and very good question has been asked already, but has no answer:</p><aside class="quote quote-modified" data-post="21" data-topic="94">
  <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/bpetruzzo/48/5075_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/why-cant-i-see-grad-of-an-intermediate-variable/94/21">Why cant I see .grad of an intermediate variable?</a> 
  </div>
  <blockquote>
    Interesting stuff! I have 3 questions in connection (my first ever post here, so I beg for forgiveness about my greenness and the lengthy post <img width="20" height="20" src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title="slight_smile" alt="slight_smile" class="emoji"> 
I’m trying to implement <a href="https://arxiv.org/abs/1711.02257" rel="noopener nofollow ugc">GradNorm</a>, a strategy to dynamically adapt weights for the individual loss contributions in a multi-task learning scenario. So, similar to your example and to the <a href="https://discuss.pytorch.org/t/how-to-combine-two-gradients-in-multi-task-model/11941/4">post</a> of <a class="mention" href="/u/hubert0527">@hubert0527</a> think: 
parent module C, 
“ending” in some last common conv-layer having weights/bias {W}, feeding into individual 
branch modules A an…
  </blockquote>
</aside>

<p>Is there a way to efficiently perform derivatives w.r.t. other losses without retain_graph?</p>
<ul>
<li>e.g., calling backward() on copies of W (W1,W2)  w.r.t L1, then L2…</li>
<li>e.g., multidimensional loss L = [l1,l2,lTOT] and then L.backward(Tensor) with Tensor somehow pointing to the right loss [1,0,0], [0,1,0],[0,0,1]</li>
</ul>
<p>Additionaly,  would it be possible to clarify how to inpsect and visualize which are the tensor responsible for memory leaks?</p>
<p><a class="mention" href="/u/alband">@albanD</a>  <span class="mention">@colensbury</span> could you help? your opinions are very valuable</p>
<p>Many thanks for your help,<br>
Stefano</p>

albanD: <p>Hi,</p>
<p>All the losses seem to share the same computation graph. So to be able to use that graph multiple times when you do multiple backward calls, you have to keep it around with <code>retain_graph=True</code>. I don’t see any way around this but reduce the memory usage by reducing batchsize.</p>

Stefano_Savian: <p>thanks so much for your prompt reply!<br>
I understand what you mean. So this is one of the cases where “retain_graph = True” cannot be avoided.<br>
However for me this does not work. Any batchsize at a certain point of the training (few to many k iterations leads to GPU OOM, for a model which requires 3 Gb on average without retain graph.</p>
<p>I have tried to avoid all memory leaks, and still the GPU is full when using retain graph. Why does the network grow so much in size? and why does it keep track of all previous iterations of the network? if I can ask?</p>
<p>Additional question: what is the recommended way to analyze tensor and graph memory usage for debugging?</p>
<p>Thanks a lot!</p>

albanD: <p>The memory leak is most likely not related to retain_graph but other changes you make. The only different retain_graph makes is that it delays the deletion of some buffers until the graph is deleted.<br>
So the only way to these to leak is if you never delete the graph. But if you never delete it, even without retain_graph, you would end up running out of memory.</p>
<p>I would check using a package like <a href="https://github.com/szagoruyko/pytorchviz">torchviz</a> to make sure the graph does not grow forever. If it does, then you need to fix your code to avoid that (most likely missing .detach()).</p>

Stefano_Savian: <p>amazing! thanks for your help!</p>

bpetruzzo: <p>Hi Stefano,</p>
<p>Sorry for the delayed reply!</p>
<p>It was a long time ago that I experimented with this and I haven’t pursued the issue much further but here is a workaround that helped for me, but this is for a 2 tasks example only (I wouldn’t speak of memory “leak” by the way, retain_graph=True just requires lots of memory.)</p>
<p>So suppose your total loss is:</p>
<p>loss = w_1 * loss_1 + w_2 * loss_2</p>
<pre><code class="lang-auto"># step A: calculate gradient (and its magnitude) of FIRST loss wrt a certain layer "manually", WITH retain_graph=True

gw_1 = torch.autograd.grad(
    loss_1,
    (
        myNet.&lt;whichever layer is the last common one&gt;.conv.weight,
        myNet.&lt;whichever layer is the last common one&gt;.conv.bias,
    ),
    retain_graph=True,
)

G_W_1 = torch.sqrt(torch.sum(gw_1[0] ** 2) + torch.sum(gw_1[1] ** 2))

# step B: BACKWARD through ENTIRE network
optimizer.zero_grad()
loss.backward()
optimizer.step()

# step C: now calculate gradient of SECOND loss by simply SUBTRACTING the FIRST gradient
# from the SUM that was calculated via the backward on the TOTAL loss
# (alternative would be another autograd.grad on loss_2 with retain_graph=True followed by the loss.backward,
# but that's where memory problems started for me)

</code></pre>
<pre><code class="lang-auto">gw_2 = []
gw_2.append(myNet.&lt;whichever layer is the last common one&gt;.conv.weight.grad - gw_1[0])
gw_2.append(myNet.&lt;whichever layer is the last common one&gt;.conv.bias.grad - gw_1[1])

G_W_2 = torch.sqrt(torch.sum(gw_2[0] ** 2) + torch.sum(gw_2[1] ** 2))

G_W_mean = 0.5 * (G_W_1 + G_W_2)
G_W_mean_no_grad = G_W_mean.detach()
</code></pre>
<p>Then you continue to work with these two quantities, G_W_mean, G_W_1 and G_W_2 to follow the loss-weight update strategy for w_1 and w_2 as described in the paper.</p>
<p>I guess this only works for 2 task problems though.</p>
<p>Anyway, I hope it helped a little.</p>
<p>Cheers,</p>
<p>Bernie</p>

Stefano_Savian: <p>Dear <a class="mention" href="/u/alband">@albanD</a> , after experimenting a bit I still have some doubts:</p>
<ul>
<li>I have used torchviz as suggested but in my case it is not helping. I have 2 “models”: one proper model which is my network, and one two dimensional nn.Parameter to store the learnt losses weights.</li>
</ul>
<p>You suggested to check if the graph is growing in size with torchviz, but model is working properly and not growing in size. When inspecting the nn.Parameters, torchviz only shows one box with none.</p>
<p>However it seems that torchviz is of no help for visualizing the autograd backward graph. The only method I found online is get_grad(), but it is verbose.<br>
<aside class="onebox stackexchange">
  <header class="source">
      <a href="https://stackoverflow.com/questions/52988876/how-can-i-visualize-what-happens-during-loss-backward" target="_blank" rel="nofollow noopener">stackoverflow.com</a>
  </header>
  <article class="onebox-body">
      <a href="https://stackoverflow.com/users/9316652/tom" target="_blank" rel="nofollow noopener">
    <img alt="Tom" src="https://www.gravatar.com/avatar/916427eb04cc8e11a432f7bc3c679de4?s=128&amp;d=identicon&amp;r=PG&amp;f=1" class="thumbnail onebox-avatar" width="60" height="60">
  </a>
<h4>
  <a href="https://stackoverflow.com/questions/52988876/how-can-i-visualize-what-happens-during-loss-backward" target="_blank" rel="nofollow noopener">How can I visualize what happens during loss.backward()?</a>
</h4>

<div class="tags">
  <strong>pytorch, backpropagation</strong>
</div>

<div class="date">
  asked by
  
  <a href="https://stackoverflow.com/users/9316652/tom" target="_blank" rel="nofollow noopener">
    Tom
  </a>
  on <a href="https://stackoverflow.com/questions/52988876/how-can-i-visualize-what-happens-during-loss-backward" target="_blank" rel="nofollow noopener">12:05PM - 25 Oct 18 UTC</a>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>How can I visualize and understand the backward graph to understand what is generating out of memory? It seems very hard to inspect what it going on under the hood, mostly in unconventional applications such as gradNorm.</p>
<p>Many thanks,<br>
Stefano</p>

albanD: <p>Hi,</p>
<p>graphviz shows you exactly the backward graph.<br>
<code>nn.Parameter</code> are always leaf Tensors (with no history), so it is expected that you get a single box with None when plotting their history as they never have any history.</p>