Help for finding memory leak

jimmy: <p>Hi,<br>
I tried something a bit unsual. It is working, but sadly there is a memory leak. Maybe someone can help me finding it. Basicially I tried to create a gradient based integration function:<br>
Fit a model with the gradient values of a function. Here I am fitting both, the y values and the gradients in parallel.</p>
<pre><code class="lang-python">import torch
from torch import nn, optim
from torch.nn import functional as F
from torch import FloatTensor as FT


def parabel(a, b, c, x):
    return a*x**2 + b*x + c


class ParabelModel(nn.Module):
    def __init__(self, start_params=[1, 1, 0]):
        super().__init__()
        self.params = nn.Parameter(FT(start_params))
        
    def forward(self, x):
        a, b, c = self.params
        return parabel(a, b, c, x)


def print_usage(device):
    if device.type == 'cuda':
        actual_mem = torch.cuda.memory_allocated() 
        max_mem = torch.cuda.max_memory_allocated()
        print(f"{actual_mem/max_mem :0.2f} actual: {actual_mem:.2e} max: {max_mem:.2e}")
    return


a, b, c = -2, 3, 2
N = 100

device = torch.device('cuda')
# create data
xs_train = torch.linspace(-1, 1, N, requires_grad=True, device=device)
ys_train = parabel(a, b, c, xs_train)
ys_train.backward(torch.ones_like(ys_train))
dx_train = xs_train.grad
xs_train = xs_train.detach()
ys_train = ys_train.detach()

l1 = 1
l2 = 0.00
lr = 0.5

model = ParabelModel([2, 1, 0]).to(device)
optimizer = optim.AdamW(model.parameters(), lr=lr)


for epoch in range(400):
    xs = xs_train.clone().detach().requires_grad_(True)
    ys_pred = model(xs)

   # calculate gradients of the model wrt the inputs
    ys_pred.backward(torch.ones_like(ys_pred), create_graph=True)
    dx_pred = xs.grad

    loss_grad = l1 * F.mse_loss(dx_pred, dx_train)

    model.zero_grad()

    loss_grad.backward()

    loss = l2 * F.mse_loss(ys_pred, ys_train)
    loss.backward()
    
    with torch.no_grad():
        optimizer.step()
        optimizer.zero_grad()
        
    
    if not epoch % 50:
        print_usage(device)
        print(model.params.data)
</code></pre>
<p>This gives:<br>
0.63 actual: 6.14e+03 max: 9.73e+03<br>
tensor([1.4900, 1.4950, 0.0000], device=‘cuda:0’)<br>
0.41 actual: 5.73e+04 max: 1.39e+05<br>
tensor([-2.1686,  3.0714,  0.0000], device=‘cuda:0’)<br>
0.57 actual: 1.09e+05 max: 1.90e+05<br>
tensor([-1.9972,  2.9907,  0.0000], device=‘cuda:0’)<br>
0.66 actual: 1.60e+05 max: 2.41e+05<br>
tensor([-1.9885,  2.9921,  0.0000], device=‘cuda:0’)<br>
0.72 actual: 2.11e+05 max: 2.92e+05<br>
tensor([-1.9894,  2.9928,  0.0000], device=‘cuda:0’)<br>
0.76 actual: 2.62e+05 max: 3.44e+05<br>
tensor([-1.9906,  2.9937,  0.0000], device=‘cuda:0’)<br>
0.79 actual: 3.13e+05 max: 3.95e+05<br>
tensor([-1.9915,  2.9943,  0.0000], device=‘cuda:0’)<br>
0.82 actual: 3.65e+05 max: 4.46e+05<br>
tensor([-1.9922,  2.9947,  0.0000], device=‘cuda:0’)</p>
<p>It’s convergeng to the expected values a, b and c but the memory grows.<br>
Thanks</p>

tom: <p>Using torch.autograd.grad (which is a better fit conceptually)</p>
<pre><code class="lang-python">    dx_pred, = torch.autograd.grad(ys_pred, xs, torch.ones_like(ys_pred), create_graph=True)
</code></pre>
<p>removes the memory leak.<br>
Note that <code>torch.tensor</code> is preferred over the constructors like FloatTensor and we generally avoid .data. You don’t need the <code>no_grad</code> for calling opt.step.</p>
<p>That said, using backward that way probably shouldn’t leak memory. <a class="mention" href="/u/alband">@albanD</a> might know if there is a story to it.</p>
<p>Best regards</p>
<p>Thomas</p>

jimmy: <p>Hi,<br>
I really thank you a lot for your answer. It fixes my problem <img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=9" title=":grinning:" class="emoji" alt=":grinning:"> But I have to say the manipulation of the dynamic graph is still complicated for me. Do you know of a good (maybe visual) explanation of that topic? Something like a correspondance between the pytorch commands and what is happening to the graph.</p>

tom: <p>The most detailed account is in <a class="mention" href="/u/ptrblck">@ptrblck</a> and my imaginary book which focuses on PyTorch details like that (and tries to avoid explaining deep learning).<br>
The other day I wrote an extension that shows the autograd graph in Jupyter at the end of each cell (or when you want) to create the visuals and facilitate interactive exploration. But didn’t think to include .grad and I’ll admit that I don’t think it’ll say “memory leak”, so it would be an imperfect solution for the thing here.<br>
If the present book is any indication, it’ll be there in a matter of years. Maybe we should have an early access component. <img src="https://discuss.pytorch.org/images/emoji/apple/stuck_out_tongue_winking_eye.png?v=9" title=":stuck_out_tongue_winking_eye:" class="emoji" alt=":stuck_out_tongue_winking_eye:"></p>

albanD: <p>Hey,</p>
<p>Yes we added a note to the doc about this: <a href="https://pytorch.org/docs/master/autograd.html#torch.autograd.backward">https://pytorch.org/docs/master/autograd.html#torch.autograd.backward</a></p>
<p>The short answer is:</p>
<ul>
<li>If possible, use autograd.grad, it will avoir this, issues with accumulating unrelated gradients in the .grad field if you forget to .zero_grad and avoid computing un-needed gradients if you only want them for the input.</li>
<li>If you cannot, replace <code>model.zero_grad()</code> after the backward that create_graph=True by</li>
</ul>
<pre><code class="lang-auto">for p in model.parameters():
    p.grad = None
</code></pre>
<p>This will be as fast but will properly remove the leak <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

jimmy: <p>Thank you aswell. I learned a lot today.</p>

tom: <p>Thank you Alban!</p>
<p>It’s not the parameters here, btw. Here is what happens as far as I can tell:</p>
<ul>
<li>Obviously, <code>xs</code> holds a reference to <code>xs.grad</code> if that is populated, as it is a member.</li>
<li>But via the graph built with <code>backward(..., create_graph=True)</code> you have that <code>xs.grad</code> via its <code>grad_fn</code> and some <code>next_functions</code> has an indirect reference to <code>xs</code>.</li>
</ul>
<p>This reference cycle isn’t just in Python but rather in C++, so while the Python bits will go out of scope and can be garbage collected eventually, the C++ doesn’t have garbage collection and so the C++ side of the tensors will remain alive.<br>
If you set <code>xs.grad</code> to  <code>None</code>, it’ll remove the C++ link and so the C++ bits will be deallocated as well.</p>
<p>Best regards</p>
<p>Thomas</p>