Memory leak in LibTorch, extremely simple code

Ben_Harper: <p>I have the following tiny code snippet, which allocates a new chunk of CUDA memory every time I call model.forward():</p>
<pre><code class="lang-cpp">auto model = torch::jit::load("model.tm");
auto out   = torch::empty({1, 512, 512});
for (int i = 0; i &lt; 20; i++) {
	auto in  = torch::empty({1, 3, 512, 512}, torch::kCUDA);
	auto res = model-&gt;forward({in}).toTensor(); // every time this is called, another 2GB CUDA memory
	out.copy_(res[0]); // without this line, no leaks
}
</code></pre>
<p>Obviously, this runs out of CUDA memory very quickly.<br>
Curiously, if I don’t consume the result of model.forward(), then no leaks.<br>
Am I doing something wrong?</p>
<p>I’m on Ubuntu 18.04, CUDA 10.0, and PyTorch compiled from source, v1.0.1 branch.</p>

Ben_Harper: <p>I have found a workaround (pseudo-code):</p>
<pre><code class="lang-cpp">float* buf = (float*) malloc(bufSize);
for (int i = 0; i &lt; 20; i++) {
    // input is a CUDA tensor
    auto res = model-&gt;forward(input);
    auto resCPU = res.cpu();
    memcpy(buf, resCPU.accessor&lt;float, 3&gt;().data(), bufSize);
    // resCPU goes out of scope here, so presumably some or other reference count
    // is also release, so that the CUDA memory involved in the model.forward() can now
    // be released.
}
</code></pre>
<p>The key difference here is that there is no long-lived object that is the destination of a CUDA -&gt; CPU copy. In my original example, the ‘out’ object seems to cause references to the CUDA memory to never be released.<br>
Is this intended behaviour?</p>

albanD: <p>Hi,</p>
<p>The problem is that you’re doing a forward in a network. And if you don’t disable the autograd, then all the buffers are saved to be able to backprop when you want.<br>
Then you save that output in another Tensor thus keeping this whole history around.</p>
<p>If you just want to do evaluation, you should disable <code>GradMode</code> (not sure where it is in the cpp doc).</p>

shane-carroll: <p>The c++ equivalent of <code>torch.no_grad()</code> would be <code>NoGradGuard</code> from<br>
<a href="https://github.com/pytorch/pytorch/blob/master/torch/csrc/api/include/torch/utils.h" rel="nofollow noopener">torch/csrc/api/include/torch/utils.h</a>. From the current comments you can see that it is a thread-local guard to disable gradients.</p>
<p>torch::NoGradGuard no_grad_guard;<br>
[rest of your code with no grad]</p>

Ben_Harper: <p>Thanks <a class="mention" href="/u/alband">@albanD</a>, <a class="mention" href="/u/shane-carroll">@shane-carroll</a>!</p>
<p><code>torch::NoGradGuard no_grad_guard</code> works!</p>
<p>Memory usage is also over 2x less, which makes sense.</p>
<p>This really ought to be in the “getting started” docs. I’ll try and find time to make a little PR that includes this in the documentation.</p>

Z_Froyo: <p>when I am on CPU, I have the same problem. I tried [torch::NoGradGuard no_grad_guard;] but its of no use. I compile the code to dll for another program to use. There’s no problem when I run my code independently. In another program, its still ok when the function [forward] is annotated…I really dont know how to solve, anyone knows?please help thanks!</p>

Kai_Huang: <p>I also encountered this question. Using <code>torch::NoGradGuard no_grad_guard</code> works.<br>
Want to ask what if I want to backward as well? In that case we can’t set NoGradGuard…<br>
Thanks for the help in advance!</p>

albanD: <p>If you want to backward, then you need to keep these buffers. Be careful though only to link things that make sense for your problem.<br>
If you want to accumulate the loss for example, you will need to use detach() to make sure the history won’t be kept.</p>

Kai_Huang: <p>Thanks so much for your reply!</p>
<p>So after accumulating loss or getting gradients, what variables need to be detached?</p>

albanD: <p>The variables that naturally go out of scope don’t need any special treatment.<br>
If you keep some things around across iterations, for example if you accumulate the loss for every batch of the epoch, then you want to detach before accumulating in this buffer.</p>

Kai_Huang: <p>Oh, I see. What if I just get the loss, gradient input and weights, do I need to do something special?<br>
Thanks!</p>

albanD: <p>That depends what you do with them <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
If you return that to other code as plain values. Then you should detach them before returning them.<br>
If you just use them locally to update the weights, nothing is needed.</p>