GPU memory leak

username: iftg
cooked: <p>I run out of GPU memory when training my model. The leak seems to be happening at the first call of loss.backward(). I guess that somehow a copy of the graph remain in the memory but can’t see where it happens and what to do about it.</p>
<p>Here’s my fit function:</p>
<pre><code class="lang-auto">    val_loss_best = np.inf
    losses = []
    
    # Prepare loss history
    for epoch in range(epochs): 
        for idx_batch, (x, y) in enumerate(dataloader_train):            
        
            optimizer.zero_grad()
            
            # Propagate input
            mem_gpu = torch.cuda.memory_stats(device='cuda:0')["allocated_bytes.all.current"]
            print(f'Before forward: {mem_gpu:,}')
                            
            netout = net(x.to(device), y.to(device))
            mem_gpu = torch.cuda.memory_stats(device='cuda:0')["allocated_bytes.all.current"]
            print(f"After forward: {mem_gpu:,}")

            # Compute loss on the training set
            loss = loss_function(netout, y.to(device))
            mem_gpu = torch.cuda.memory_stats(device='cuda:0')["allocated_bytes.all.current"]
            print(f'After loss calc: {mem_gpu:,}')
                
            # Backpropage loss
            loss.backward()
            mem_gpu = torch.cuda.memory_stats(device='cuda:0')["allocated_bytes.all.current"]
            print(f'After backward: {mem_gpu:,}')

            # Update weights
            optimizer.step()
  
        # Compute loss on the validation set   
        net.eval()
        mem_gpu = torch.cuda.memory_stats(device='cuda:0')["allocated_bytes.all.current"]
        print(f'After net.eval(): {mem_gpu:,}')
        
        # mem_gpu = torch.cuda.memory_stats(device='cuda:0')["allocated_bytes.all.current"]
        # print(f'After torch.no_grad():  {mem_gpu:,}')
            
        val_loss = compute_loss(net, dataloader_val, loss_function, device) #float
        mem_gpu = torch.cuda.memory_stats(device='cuda:0')["allocated_bytes.all.current"]
        print(f'After val loss calc: {mem_gpu:,}')
        
        losses.append(val_loss)
        
        if val_loss &lt; val_loss_best:
            val_loss_best = val_loss
        
        net.train()
          
    return val_loss_best
</code></pre>
<pre><code class="lang-auto">def compute_loss(net: torch.nn.Module,
                 dataloader: torch.utils.data.DataLoader,
                 loss_function: torch.nn.Module,
                 device: torch.device = 'cpu') -&gt; float:

    running_loss = 0
    with torch.no_grad():
        for idx_batch, (x, y) in enumerate(dataloader): #iterate across batches
            netout = net(x.to(device))
            current_loss = loss_function(y.to(device), netout).item()
            running_loss += current_loss

    return running_loss / len(dataloader)
</code></pre>
<p>and loss_function is nn.MSE().<br>
The output of this code is:<br>
Before I call fit():<br>
Before initialising model: 534,819,328<br>
After initialising model: 635,895,808</p>
<p>Now we go into fit():</p>
<p>Before forward: 635,895,808<br>
After forward: 647,901,184<br>
After loss calc: 647,902,208<br>
After backward: 711,877,120</p>
<p>Before forward: 787,703,808<br>
After forward: 799,555,072<br>
After loss calc: 799,555,584<br>
After backward: 863,530,496</p>
<p>Before forward: 787,703,808<br>
After forward: 799,555,072<br>
After loss calc: 799,555,584<br>
After backward: 863,530,496<br>
… (these four lines will repeat until we exit from fit()).<br>
Then, After moving model to cpu: 686,627,328</p>
<p>So I’m losing 150MB of GPU memory. Any ideas, please help!</p>

username: ptrblck
cooked: <p>Which optimizer are you using? If it’s holing internal states an increase in memory is expected after the first <code>step()</code> call and you are not facing a memory leak.</p>

username: iftg
cooked: <p>Adam.<br>
Actually, on a more careful glance, the leak might be happening on the second call (as well?).<br>
I am not sure what you mean by “not facing a memory leak”. I finish my training having 150MB less than when I start it and can’t see what consumes this memory. Since I am doing this training as part of a loop (models change over time) I exhaust my GPU memory pretty quickly.<br>
In other words, no matter whether or not it’s a leak there must be a way to release this memory. I just don’t see how.<br>
What am I missing?</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="iftg" data-post="3" data-topic="193572">
<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/i/ecccb3/48.png" class="avatar"> iftg:</div>
<blockquote>
<p>I am not sure what you mean by “not facing a memory leak”.</p>
</blockquote>
</aside>
<p>The memory increases as expected and saturates at 863MB, which does not indicate any leaks.</p>
<aside class="quote no-group" data-username="iftg" data-post="3" data-topic="193572">
<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/i/ecccb3/48.png" class="avatar"> iftg:</div>
<blockquote>
<p>I finish my training having 150MB less than when I start it and can’t see what consumes this memory.</p>
</blockquote>
</aside>
<p>Which is expected, since you are using <code>Adam</code>, which will create internal states for each param, and are not checking the memory after the <code>optimizer.step()</code> call.</p>
<aside class="quote no-group" data-username="iftg" data-post="3" data-topic="193572">
<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/i/ecccb3/48.png" class="avatar"> iftg:</div>
<blockquote>
<p>Since I am doing this training as part of a loop (models change over time) I exhaust my GPU memory pretty quickly.</p>
</blockquote>
</aside>
<p>I don’t understand how this can happen, if the memory usage is static based on your comment:</p>
<aside class="quote no-group" data-username="iftg" data-post="1" data-topic="193572">
<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/i/ecccb3/48.png" class="avatar"> iftg:</div>
<blockquote>
<p>Before forward: 787,703,808<br>
After forward: 799,555,072<br>
After loss calc: 799,555,584<br>
After backward: 863,530,496<br>
… (these four lines will repeat until we exit from fit()).</p>
</blockquote>
</aside>
<aside class="quote no-group" data-username="iftg" data-post="3" data-topic="193572">
<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/i/ecccb3/48.png" class="avatar"> iftg:</div>
<blockquote>
<p>In other words, no matter whether or not it’s a leak there must be a way to release this memory. I just don’t see how.</p>
</blockquote>
</aside>
<p>No, since <code>Adam</code> requires this memory. You would need to delete the optimizer to be able to release the memory.</p>

username: iftg
cooked: <p>The memory usage is static in my comment because it’s just one training of the model (at one point in time). These numbers are produced within loop</p>
<pre><code class="lang-auto">for epoch in range(epochs): 
        for idx_batch, (x, y) in enumerate(dataloader_train):
</code></pre>
<p>in my code. Then I move to the next data point (the next point in time, which mean a new window and new dataloader_train and dataloader_val) and train another model, which will consume another 150MB.<br>
It had not occurred to me that the problem might be caused by an optimizer.</p>
<pre><code class="lang-auto">del optimizer
torch.cuda.empty_cache()
</code></pre>
<p>seems to do the trick but there’s still about 20MB left after the first time step and a couple of MB which adds at every ensueing time step. What can it be and what to do about it?</p>

username: iftg
cooked: <p>I have identified the problem. It turns out that I had an assignment to a tensor, which was a class attribute, in the forward pass, something like:</p>
<pre><code class="lang-auto">self._ten = torch.bmm(...)
</code></pre>
<p>It was enough to change it to:</p>
<pre><code class="lang-auto">ten = torch.bmm(...)
</code></pre>