Load_state_dict causes memory leak

username: grahammorehead
cooked: <p>If you store a state_dict using torch.save, and then load that state_dict (or another), it doesn’t just replace the weights in your current model.  It loads the new values into GPU memory and then <em>maybe</em> releases the old GPU memory.</p>
<p>If I set my vector length to 4900, PyTorch eventually releases unused GPU memory and everything goes fine…  If I set it to 5000, however, GPU memory usage increases and it never lets go.</p>
<p>To reproduce:</p>
<p>(First some info:)<br>
Hardware: AWS P3.2xlarge (which has a Tesla V100)</p>
<blockquote>
<p>uname -a<br>
Linux ip-10-5-5-70 4.4.0-1074-aws <span class="hashtag-raw">#84-Ubuntu</span> SMP Thu Dec 6 08:57:58 UTC 2018 x86_64 x86_64 x86_64 GNU/Linux</p>
</blockquote>
<blockquote>
<p>torch.<strong>version</strong><br>
‘1.0.0’</p>
</blockquote>
<pre><code class="lang-auto">
import sys
import os
import time
import torch
import torch.nn as nn
import torch.nn.functional as F

######################
def get_GPU_memory():
    """
    Return a list of how much memory is currently used by each GPU
    """
    try:
        dev = int(os.environ['CUDA_VISIBLE_DEVICES'])
    except:
        dev = 0
    mem = torch.cuda.memory_allocated(device=dev)
    return mem

#################################
class StatePredictor(nn.Module):
    def __init__(self, dim, options={}):
        super(StatePredictor, self).__init__()
        self.relu     = torch.nn.LeakyReLU(negative_slope=0.001, inplace=False)
        self.norm      = nn.LayerNorm(dim, elementwise_affine=False)
        self.fc1       = nn.Linear(dim, dim).double()
        self.fc2       = nn.Linear(dim, dim).double()
        self.fc3       = nn.Linear(dim, dim).double()
        self.fc4       = nn.Linear(dim, dim).double()
        self.fc5       = nn.Linear(dim, dim).double()
        self.fc6       = nn.Linear(dim, dim).double()
        self.fc_final  = nn.Linear(dim, 2).double()
        self.dropout_p = 0.1


    def forward(self, H):
        x = H
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = torch.tanh(x)
        x = self.fc3(x)
        x = self.relu(x)
        x = self.fc4(x)
        x = torch.tanh(x)
        x = self.fc5(x)
        x = self.relu(x)
        x = self.fc6(x)
        x = torch.tanh(x)
        x = self.fc_final(x)
        x = F.softmax(x, dim=0)
        return x

    
###########################
if __name__ == '__main__':

    print("\n[%d] MEMORY before creating layers\n"% get_GPU_memory())
    dirpath   = "models"
    dim       = int(sys.argv[1])
    predictor = StatePredictor(dim)                                                   # Create a simple fully-connected network
    if torch.cuda.is_available():
        predictor = predictor.cuda()                                                  # Load the model into the GPU
        
    print("[%d] MEMORY after creating layers\n"% get_GPU_memory())
    torch.save(predictor.state_dict(), dirpath+'/StatePredictor.pth')                 # Save the state_dict

    # If 'load_state_dict' replaces the model's weights and biases, then why does GPU memory usage go up each time?
    step = 0
    while True:
        step += 1
        predictor.load_state_dict(torch.load(dirpath+'/StatePredictor.pth'))         # Load state_dict
        if step &gt; 20:
            time.sleep(10)
        torch.cuda.empty_cache()        
        print("%d.) [%d] MEMORY after loading state_dict"% (step, get_GPU_memory()))

###########################
</code></pre>
<p>Usage:</p>
<blockquote>
<p>python test.py 5000</p>
</blockquote>

username: JuanFMontesinos
cooked: <p>I recommend you to load the state dict in cpu and then to move it to gpu.</p>

username: grahammorehead
cooked: <p>I’ll try that now and then post results here.</p>

username: grahammorehead
cooked: <p><a class="mention" href="/u/juanfmontesinos">@JuanFMontesinos</a>,</p>
<p>I added two lines around the ‘load_state_dict’ line:</p>
<pre><code class="lang-auto">        predictor = predictor.cpu()
        predictor.load_state_dict(torch.load(dirpath+'/StatePredictor.pth'))         # Load state_dict
        predictor = predictor.cuda()
</code></pre>
<p>It barely helps.  Instead of being able to use a max vector length of 4900, now it’s 5000.  The problems are pushed out to 5100.</p>

username: JuanFMontesinos
cooked: <p>The problem is, (talkin by experience), is that we save cuda weights. However, when we load the state dict, we open those weights on the gpu and memory is never freed.</p>
<p>If i’m not wrong, opening weights this way should work (at least talking about gpu vram, never checked cpu ram):</p>
<pre><code class="lang-auto">state_dict = torch.load(directory, map_location=lambda storage, loc: storage)
</code></pre>
<p>so then you can load that state dict</p>

username: grahammorehead
cooked: <p><a class="mention" href="/u/juanfmontesinos">@JuanFMontesinos</a>, Thank You!</p>
<p>I put it like this:</p>
<pre><code class="lang-auto">        state_dict = torch.load(dirpath+'/StatePredictor.pth', map_location=lambda storage, loc: storage)
        predictor.load_state_dict(state_dict)
</code></pre>
<p>Does that look right?</p>
<p>It certainly doesn’t have a memory leak now.</p>

username: JuanFMontesinos
cooked: <p>seems fine for me, however, as i said i never check cpu ram <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=6" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: grahammorehead
cooked: <p>Who cares – that’s not where the magic happens <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=6" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: JuanFMontesinos
cooked: <p>Ohhh you are not woking wih video right? I get out of ram at the time of loading data hahahah</p>

username: nuoyi_ma
cooked: <p>nice，thank you<br>
it is a prefect way to deal with it</p>

username: Eta_C
cooked: <p><strong>DELETE STATE DICT</strong> work for me</p>
<pre><code class="lang-python">state_dict = torch.load(state_dict_path) 
model.load_state_dict(state_dict)
del state_dict
</code></pre>
<p>But I don’t know if the weights in the model would be deleted after <code>del state_dict</code>. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: SimonW
cooked: <p>they wont be deleted.</p>

username: jxgu1016
cooked: <p>hi guys, check out this <a href="https://github.com/pytorch/pytorch/issues/7415#issuecomment-693424574" rel="nofollow noopener">solution</a>. It works for me.</p>

username: jxgu1016
cooked: <p>btw,</p>
<pre><code class="lang-auto">del checkpoint
torch.cuda.empty_cache()
</code></pre>
<p>also helps, if you gpu can pass the memory peak.</p>