GPU memory usage steadily increases in forward pass

username: MariosOreo
cooked: <p>Hi,</p>
<p>I check GPU memory usage by <code>torch.cuda.max_memory_allocated()</code>, and it steadily increases in each forward pass. In the forward pass, it needs to store two intermediate tensors to calculate the next result like Fibonacci. So I think the problem may be in the way that I access these intermediate tensors.<br>
My method:</p>
<pre><code class="lang-auto">def forward(self, x):
    ...
    intermediate_result = [] # a tensor list with length of 2, store results of i-1 and i-2
    for i in range(nb_loops):
        next_result = func(intermediate_result) # func is used to calculate next_result
        intermediate_result.pop()
        intermediate_result.append(next_result)
        ...
</code></pre>
<p>Could you give me some advice on how to maintain such intermediate_result in forward pass, and how to avoid GPU memory usage increasing problem?<br>
Thanks in advance.</p>

username: albanD
cooked: <p>Hi,</p>
<p>This looks like a good way to maintain them.<br>
What you need to make sure is how big your computational graph is. In particular, if <code>intermediate_result</code> requires gradients, then <code>next_result</code> will also <strong>and</strong> will keep track of its history (which contains <code>intermediate_result</code>).</p>

username: MariosOreo
cooked: <p>Thanks for your reply!</p>
<p>What you mean is that if <code>intermediate_result</code> requires gradients, it will not only maintain two tensors (in <code>intermediate_result</code> currently), and it will also track all the preceding tensors in the computational graph, right? If so, the steadily growing memory usage the a normal case.</p>
<aside class="quote no-group" data-username="albanD" data-post="2" data-topic="61422">
<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>What you need to make sure is how big your computational graph is.</p>
</blockquote>
</aside>
<p>As you said, the memory usage increases in the first epoch, and it becomes stable in the following phase.</p>

username: albanD
cooked: <p>Yes it will track all the previous ones.</p>
<p>Can you explain which part you want to backpropagate through?</p>
<ul>
<li>Do you want to go all the way to the first iteration?</li>
<li>Do you want to backprop only the last <code>func</code> eval?</li>
<li>Do you want to backprop only the last two?</li>
</ul>

username: MariosOreo
cooked: <p>Sorry for late reply.</p>
<p>Yes, I want to go all the way to the first iteration, backprop to <code>i_0</code> (i.e. input of the network). Additionally, during forward pass, in each iteration, the selection of intermediate feature <code>i_k</code> (<code>i_k</code> can have different size, that means it will not have a constant GPU memory consumption) based on Gumbel-Softmax, which also consumes additional GPU memory (I think).</p>
<p>So, in the forward pass, GPU memory checked by <code>nvidia-smi</code>, will sometimes increase, sometimes be stable, and it can also raise OOM anytime (maybe forward with large intermediate features?)</p>
<p>I want to know, how can I control the GPU memory usage? Does Gumbel-softmax of each iteration cause such growing memory usage? Could I use <code>torch.cuda.empty_cache()</code> after each epoch to release memory usage?</p>
<p>Thanks for you help!</p>

username: albanD
cooked: <p>Hi,</p>
<p>If you need gradients all the way back to the first iteration, it is expected that the memory usage is going to be quite high as it needs to keep track of all the buffers to be able to compute the gradients.</p>
<blockquote>
<p>Does Gumbel-softmax of each iteration cause such growing memory usage?</p>
</blockquote>
<p>It is all the intermediary results that are required to compute the gradients that are expensive. You can try to use <a href="https://pytorch.org/docs/stable/checkpoint.html"><code>checkpoint</code></a> to reduce the memory usage at the cost of more compute.</p>
<blockquote>
<p>Could I use  <code>torch.cuda.empty_cache()</code>  after each epoch to release memory usage?</p>
</blockquote>
<p>That won’t reduce the memory usage as the memory in the cache is reused to allocate more Tensor. The only thing this will do is slow down your code by forcing emptying and refilling the cache a lot.</p>