GPU memory usage increases by ~90% after torch.load

austin: <p>I’m currently training a faster-rcnn model. Normal training consumes ~1900MiB of gpu memory. When I try to resume training from a checkpoint with torch.load, the model takes over 3000MiB. With identical settings specified in a config file.</p>
<p>I believe these are the relevant bits of code:</p>
<pre><code>voc_dataset = PascalVOC(DATA_PATH, transform, LIMIT)
voc_loader = DataLoader(voc_dataset, shuffle=SHUFFLE, pin_memory=True)
basecnn = BaseCNN(ARCHITECTURE, REQUIRES_GRAD)
rpn = RegionProposalNetwork(batch_size=RPN_BATCH_SIZE)
detector = Detector(batch_size=CLF_BATCH_SIZE)

frcnn = FasterRCNN(basecnn, rpn, detector)
optimizer = Adam(filter(lambda p: p.requires_grad, frcnn.parameters()), lr=LEARNING_RATE)

if RESUME_PATH:
    experiment = torch.load(RESUME_PATH)
    frcnn.load_state_dict(experiment['model_state'])
    optimizer.load_state_dict(experiment['optimizer_state'])
    frcnn.basecnn.finetune()

frcnn.cuda()
</code></pre>
<p>Full training code is here: <a href="https://github.com/A-Jacobson/faster_rcnn/blob/master/train.py" rel="nofollow noopener">https://github.com/A-Jacobson/faster_rcnn/blob/master/train.py</a></p>

roytseng: <p>Yes, I have experienced the same situation before. This is quiet annoying ! I will bump on OOM if I would like to resume training from a checkpoint.</p>

ptrblck: <p>Have you tried to save and load the <code>state_dict</code>? I’ve seen some issues where e.g. the training resumed when using the second approach explained <a href="http://pytorch.org/docs/master/notes/serialization.html">here</a>.</p>

roytseng: <p>Yes. I always follow the best practice to save and load the <code>state_dict</code>.<br>
I found a related issue <a href="https://github.com/facebookresearch/fairseq-py/issues/51" rel="nofollow noopener">here</a><br>
It says <code>torch.cuda.empty_cache()</code>might help, but in my case I still have OOM.<br>
By the way, I’m using pytorch 0.3.1</p>

roytseng: <p>Good news ! I seems able to avoid OOM by the following loading strategy:</p>
<pre><code class="lang-python">checkpoint = torch.load(ckpt_file, map_location=lambda storage, loc: storage)
net.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])  # pytorch 0.3.1 has a bug on this, it's fix in master
del checkpoint  # dereference seems crucial
torch.cuda.empty_cache()
</code></pre>

ptrblck: <p>That’s good!<br>
Btw, shouldn’t the checkpoint be in CPU memory when loading with the <code>map_location</code> parameter?</p>

roytseng: <p>Yes. With <code>map_location=lambda storage, loc: storage</code>, tensors in checkpoint are in CPU memory at first.<br>
However, I guess <code>load_state_dict</code> may cast tensors to the corresponding device of model parameters internally, and the references to the casted tensors are still held by <code>checkpoint</code>. I don’t really trace it down.</p>

austin: <p>I’ll have to check this on my end. Embarrassingly, I ended up getting a new gpu and the temporary spike didn’t matter all that much anymore.</p>

zjoe: <p>Hi roytseng, I face the same problem when using 0.3.1, and it still blocking me. I have tried del+cuda.empty_cache() but it doesn’t work in my case. I notice your comment “pytorch 0.3.1 has a bug on this, it’s fix in master”, could you explain more on this? Links to the original issues or commits would be really helpful(i have checked commits about optimizers but do not find it). Thanks in advance!</p>

roytseng: <p>Hi zjoe,</p>
<p>I not really remember if the bug in load_state_dict of optimizer is related to the memory usage increment or not (I guess it’s not). However, it’s sure that this bug has been fixed in pytorch0.4.</p>

Daniel_Dagnino: <p>Hi all,<br>
Have you been able to fix the problem?<br>
I am experiencing the same problem and I am using pytorch 0.5, so it does not seem that the new versions solve the problem. I am using:</p>
<pre><code class="lang-auto">if load:
    checkpoint = torch.load('./model.ckpt')
    startEpoch = checkpoint['StartEpoch']
    model.load_state_dict(checkpoint['state_dict'])
    del checkpoint
    torch.cuda.empty_cache()
</code></pre>
<p>but the problem persists. The problem is located in the point</p>
<pre><code class="lang-auto">loss.backward
</code></pre>
<p>however, the model is the same; and I am not loading anything else than the model and the epoch number. I need to reduce the batch size which is very annoying.<br>
Thanks,<br>
Dani.</p>

kangkang_li: <p>have you solved it? I meet this too.</p>

1hada: <pre><code class="lang-auto">del checkpoint  # dereference seems crucial
</code></pre>
<p>Worked for me !<br>
It seems to have saved around +500mb in my case!<br>
I was able to save but not load in pytorch. The checkpoint definitely took up valuable valuable gpu memory.</p>
<pre><code class="lang-auto">torch.cuda.empty_cache() 
</code></pre>
<p>Helped clear around 300mb!</p>
<p>Thank you <a class="mention" href="/u/roytseng">@roytseng</a>  !</p>

realiti: <p>This thread is really old, but I have to report this can still be an issue with pytorch 1.8.1. I had to lower model’s batch size from 64 to 40 to just be able to resume from a checkpoint before for example. That’s quite a performance hit actually. This method works beautifully and is still relevant, cheers.</p>
<p>Edit: I can also confirm map_location=cpu is also works.</p>

Bled_Clement: <p>Still an issue. Ran into this today on Pytorch 1.10.</p>
<p>Deleting checkpoint and emptying cache doesn’t work.</p>
<p>Loading model to cpu first works as pointed out by <a class="mention" href="/u/realiti">@realiti</a>.</p>

DariusOneTrick: <p>Still an issue. My model takes 31000MiB, and I got a 32GB gpu. If I start training from scratch, stuff works; If I resume from a checkpoint, I get an OOM.</p>

KCLam: <p>This solution helped me.</p>
<pre><code class="lang-auto">state = torch.load(f, map_location=torch.device('cpu'))
</code></pre>
<p>from <a href="https://stackoverflow.com/questions/58060885/pytorch-gpu-memory-increase-after-load-operation" class="inline-onebox" rel="noopener nofollow ugc">Pytorch GPU memory increase after load operation - Stack Overflow</a></p>