Deleting Tensors in Context Save for Backward

eduardo4jesus: <p>Are the tensor saved for backward as below freed or deleted automatically after the backward pass?</p>
<pre><code class="lang-auto">ctx.save_for_backward(input, weight, bias)
</code></pre>
<p>I am trying to get around memory used problems.</p>

ptrblck: <p>Yes, these tensors should be freed after the <code>backward()</code>.<br>
To double check it, you could use <a href="https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html">this example</a> and add some <code>print</code> statements to check the memory:</p>
<pre><code class="lang-python">for t in range(5):
    # To apply our Function, we use Function.apply method. We alias this as 'relu'.
    relu = MyReLU.apply

    # Forward pass: compute predicted y using operations; we compute
    # ReLU using our custom autograd operation.
    print('='*10, t, '='*10)
    print(torch.cuda.memory_allocated()/1024)
    y_pred = relu(x.mm(w1)).mm(w2)
    print(torch.cuda.memory_allocated()/1024)

    # Compute and print loss
    loss = (y_pred - y).pow(2).sum()
    print(torch.cuda.memory_allocated()/1024)
    if t % 100 == 99:
        print(t, loss.item())

    # Use autograd to compute the backward pass.
    loss.backward()
    print(torch.cuda.memory_allocated()/1024)

    # Update weights using gradient descent
    with torch.no_grad():
        w1 -= learning_rate * w1.grad
        w2 -= learning_rate * w2.grad

        # Manually zero the gradients after updating weights
        w1.grad = None
        w2.grad = None
    
    print(torch.cuda.memory_allocated()/1024)
</code></pre>
<p>which shows that the memory falls down to the initial usage:</p>
<pre><code class="lang-python">========== 0 ==========
647.5
700.0
703.0
1045.5
650.5
========== 1 ==========
650.5
700.5
703.0
1045.5
650.5
========== 2 ==========
650.5
700.5
703.0
1045.5
650.5
========== 3 ==========
650.5
700.5
703.0
1045.5
650.5
========== 4 ==========
650.5
700.5
703.0
1045.5
650.5
</code></pre>
<p>Note that I’ve replaced the <code>.grad</code> attributes with <code>None</code> to free these tensors as well instead of zeroing them out.</p>

eduardo4jesus: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="122917">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Note that I’ve replaced the <code>.grad</code> attributes with <code>None</code> to free these tensors as well instead of zeroing them out.</p>
</blockquote>
</aside>
<p>Excellent, thank you so much. I had not thought yet about assigning the <code>.grad</code> to <code>None</code> after the <code>backward</code>. I totally missed that!</p>
<p>Currently, I was only deleting tensors whenever they were not needed anymore, such as in <code>X = torch.fft.fft2d(x); del x;</code>. Deleting the model weight’s <code>.grad</code> will definitely improve things!</p>
<p>Thanks a lot!</p>

ptrblck: <p>In case you are using an optimizer or are calling <code>zero_grad</code> on the module itself, note that you can use the <code>set_to_none=True</code> argument in <a href="https://pytorch.org/docs/stable/optim.html#torch.optim.Optimizer.zero_grad"><code>zero_grad</code></a> for the exact same reason of saving memory. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

eduardo4jesus: <p>Thank you so much again for these precious tips. I just had another question on this topic. Is there a way to free the tensors <em>saved for backwards</em> or the <code>grad_output</code> before the end of <code>backward</code>?</p>
<p>Say I have something like:</p>
<pre><code class="lang-auto">def backward(cls, ctx, grad_output):
.
.
.
   del grad_output;
.
.
.
</code></pre>
<p>I imagine that the above is pointless since there would still be a reference for <code>grad_output</code> on the call stack, right? Would there be a way do to something like this?</p>

ptrblck: <p>It doesn’t seem to make a difference, if I <code>del</code> the tensor and check the memory before and after this operation.</p>

eduardo4jesus: <aside class="quote no-group" data-username="ptrblck" data-post="6" data-topic="122917" data-full="true">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>It doesn’t seem to make a difference, if I <code>del</code> the tensor and check the memory before and after this operation.</p>
</blockquote>
</aside>
<p>Yes, I totally agree.</p>
<aside class="quote no-group" data-username="eduardo4jesus" data-post="5" data-topic="122917">
<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/eduardo4jesus/48/20203_2.png" class="avatar"> eduardo4jesus:</div>
<blockquote>
<p>I imagine that the above is pointless since there would still be a reference for <code>grad_output</code> on the call stack</p>
</blockquote>
</aside>
<p>As far as I know, Python is heavily based on references instead of pointers and it manages memory allocation on its own. But, I wonder if there would be away to do something inside <code>backward</code> that would invalidate the reference in the call stack, or somehow just forcing a memory deallocation. I am just realizing now this would be a OFF-(Torch)-Topic question.</p>

soulitzer: <p>Tensors saved for the backward pass from forward ARE actually freed automatically during the corresponding backward!</p>
<p>If you are curious you can check in the autograd engine’s evaluate function, we call <code>fn.release_variables</code> if <code>keep_graph</code> (i.e., the retain_graph parameter) is set to False.</p><aside class="onebox githubblob">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/engine.cpp#L769" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/autograd/engine.cpp#L769" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/master/torch/csrc/autograd/engine.cpp#L769</a></h4>



  <pre class="onebox">    <code class="lang-cpp">
      <ol class="start lines" start="759" style="counter-reset: li-counter 758 ;">
          <li>at::ThreadLocalStateGuard tls_guard(graph_task-&gt;thread_locals_);</li>
          <li>
          </li>
<li>// Switches to a function's CUDA stream (if applicable) before calling it</li>
          <li>const auto opt_parent_stream = (*func).stream(c10::DeviceType::CUDA);</li>
          <li>c10::OptionalStreamGuard parent_stream_guard{opt_parent_stream};</li>
          <li>
          </li>
<li>auto outputs = call_function(graph_task, func, inputs);</li>
          <li>
          </li>
<li>auto&amp; fn = *func;</li>
          <li>if (!graph_task-&gt;keep_graph_) {</li>
          <li class="selected">  fn.release_variables();</li>
          <li>}</li>
          <li>
          </li>
<li>int num_outputs = outputs.size();</li>
          <li>if (num_outputs == 0) { // Note: doesn't acquire the mutex</li>
          <li>  // Records leaf stream (if applicable)</li>
          <li>  // See note "Streaming backwards"</li>
          <li>  if (opt_parent_stream) {</li>
          <li>    std::lock_guard&lt;std::mutex&gt; lock(graph_task-&gt;mutex_);</li>
          <li>    graph_task-&gt;leaf_streams.emplace(*opt_parent_stream);</li>
          <li>  }</li>
      </ol>
    </code>
  </pre>


  </article>

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

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


eduardo4jesus: <p>Thanks for showing that piece of code. It’s very nice to know where that happens. But that was already pointed by <a class="mention" href="/u/ptrblck">@ptrblck</a>. Now I am wondering if there is a way to free these tensors (the ones saved. for backwards) and also (if possible) the <code>grad_output</code> somewhere inside the <code>backward</code> function.</p>
<p>One of the issues I face is that the calculations I do in the backward use a lot of memory. It is not really that much considering only a single layer. But, considering that some CNN architecture can get really deep and wide, I just run out of memory in such scenarios. I am looking for releasing those tensors right after using then, while still having other computations left.</p>
<p>Maybe my only hope for those scenarios is to be heavily based on in-place operations.</p>

Nikola_Andro: <p>I think the problem might be in the .step() since .backward() does free the used tensors. However, in my case, once I do optimizer.step() and it creates additional ~130 tensors and I have no clue where they are saved.</p>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a> Do you know what tensors are being created when performing optimizer.step() ?</p>
<p>Thanks</p>

Nikola_Andro: <p>After some time, setting gradients to <strong>None</strong> and deleting optimizer did the job.</p>
<pre><code class="lang-auto">optimizer.zero_grad(set_to_none=True)
del optimizer
torch.cuda.empty_cache()
</code></pre>

ptrblck: <p>It depends on the optimizer as some will create internal running stats for the parameters, such as <code>Adam</code> as seen here:</p>
<pre><code class="lang-python">model = nn.Linear(1, 1)
optimizer = torch.optim.Adam(model.parameters(), lr=1.)
x = torch.randn(1, 1)

print(optimizer.state_dict())
#{'state': {}, 'param_groups': [{'lr': 1.0, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 0, 'amsgrad': False, 'maximize': False, 'foreach': None, 'capturable': False, 'differentiable': False, 'fused': False, 'params': [0, 1]}]}

out = model(x)
out.mean().backward()
optimizer.step()

print(optimizer.state_dict())
# {'state': {0: {'step': tensor(1.), 'exp_avg': tensor([[-0.2066]]), 'exp_avg_sq': tensor([[0.0043]])}, 1: {'step': tensor(1.), 'exp_avg': tensor([0.1000]), 'exp_avg_sq': tensor([0.0010])}}, 'param_groups': [{'lr': 1.0, 'betas': (0.9, 0.999), 'eps': 1e-08, 'weight_decay': 0, 'amsgrad': False, 'maximize': False, 'foreach': None, 'capturable': False, 'differentiable': False, 'fused': False, 'params': [0, 1]}]}
</code></pre>

Nikola_Andro: <aside class="quote no-group" data-username="ptrblck" data-post="12" data-topic="122917">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p><code> torch.optim.</code></p>
</blockquote>
</aside>
<p>That makes sense. Thank you!</p>