How to print gradient graph

jwillette: <pre><code class="lang-auto">def maml_simulation() -&gt; None:
    x_inner = torch.tensor(2.0, requires_grad=True)
    x_outer = torch.tensor(3.0, requires_grad=True)
    theta_outer = torch.randn(1, requires_grad=True)
    print(f"before training: x inner: {x_inner} theta outer: {theta_outer}")

    loss_func = lambda x: (x - 10) ** 2

    loss = torch.tensor(0.0, requires_grad=True)
    for i in range(5):
        theta_outer = zero_grad(theta_outer)

        print(f"outer loop, theta before: {theta_outer}")
        theta_inner = torch.tensor(theta_outer.item(), requires_grad=True)
        for j in range(5):
            theta_inner = zero_grad(theta_inner)

            prediction = theta_inner * x_inner
            inner_loss = loss_func(prediction)

            grad = torch.autograd.grad(inner_loss, theta_inner)[0]
            with torch.no_grad():
                theta_inner = theta_inner - 0.01 * grad
            print(f"inner prediction: {prediction}, loss: {inner_loss}, grad: {grad}")

        prediction = theta_outer * x_outer
        loss = loss_func(prediction)
        print(loss.requires_grad, theta_outer.requires_grad)
        grad = torch.autograd.grad(loss, theta_outer)[0]
        print("outer grad: ", grad)

        with torch.no_grad():
            theta_outer = theta_outer - 0.01 * grad
        print(f"\n\nouter prediction: {prediction}, loss: {loss}, grad: {grad}")
        print(f"theta after: {theta_outer}\n\n")
</code></pre>
<p>I am going through some meta learning stuff and I want to try and follow the second derivatives of this loop to see what it looks likes and if it is doing what I think it is doing.</p>
<ol>
<li>
<p>Does this code calulate a second derivative as is stated in this paper? <a href="https://arxiv.org/abs/1703.03400" rel="nofollow noopener">https://arxiv.org/abs/1703.03400</a></p>
</li>
<li>
<p>How can I print out the graph or verify somehow that it is doing what I think it is doing?</p>
</li>
</ol>
<p>Thanks</p>

jwillette: <p>I simplified the above code into something more concise that shows what I am trying to do and also shows that it is not happening in pytorch.</p>
<p>By my hand calculation, the second derivative of this at the bottom print statement should be <code>-12.xx</code> but I am getting the first order derivative instead of the second even though I have set <code>create_graph=True</code>. Am I doing something wrong here?</p>
<pre><code class="lang-auto">def maml_simulation() -&gt; None:
    x_i = torch.tensor(3.0, requires_grad=True)
    x_j = torch.tensor(4.0, requires_grad=True)
    theta = torch.tensor(2.0, requires_grad=True)
    theta_two = theta.clone()
    loss_func = lambda x: (x - 10) ** 2
    print(f"before training: theta: {theta}")

    prediction = theta_two * x_i
    loss_one = loss_func(prediction)

    grad = torch.autograd.grad(loss_one, theta_two, create_graph=True)[0]
    theta_two -= 0.01 * grad
    print(
        f"first prediction: {prediction}, loss: {loss_one}, grad: {grad}, theta after update: {theta_two}"
    )

    prediction = theta_two * x_j
    loss = loss_func(prediction)
    grad = torch.autograd.grad(loss, theta_two)[0]
    print(f"second prediction: {prediction}, loss: {loss}, second grad: {grad}")
</code></pre>

albanD: <p>Hi,</p>
<p>You can use the <a href="https://github.com/szagoruyko/pytorchviz">torchviz</a> package to print the graph corresponding to gradient computations.</p>

jwillette: <p>Thanks I’ll look into that package to see if it helps. Do you see any problem with my second derivative above? Flipping the Boolean create_graph doesn’t change the second gradient at all like I would expect</p>

albanD: <p>I’m not sure what is the purpose of <code>theta_two</code> in your code above, why not use <code>theta</code> directly (turns out after more investigation below, this was the root of the problem, see the rest of the answer)?</p>
<p>Also if I read correctly, <code>loss_one = (theta * xi - 10)**2</code>.<br>
So <code>grad = 2 * xi * (theta * xi - 10)</code>.<br>
So the new <code>theta_two = theta - 0.01 * (2 * xi * (theta * xi - 10)) = theta - 0.02 * theta * xi**2 + 0.2 * xi = theta * (1 - 0.02 * xi**2) + 0.2 * xi</code>.<br>
And <code>loss = ((theta * (1 - 0.02 * xi**2) + 0.2 * xi) * xj - 10)**2 = 1.0816</code><br>
And its derivative <code>grad = 2 * xj * (1 - 0.02 * xi**2) * ((theta * (1 - 0.02 * xi**2) + 0.2 * xi) * xj - 10)</code><br>
So the final grad should be <code>-6.822399999999995</code>.</p>
<p>Why do you only see the part that correspond to the last loss computation and your code behaves the same for <code>create_graph=True</code> or <code>create_graph=False</code>?<br>
Because here: <code>grad = torch.autograd.grad(loss, theta_two)[0]</code> you ask for gradients wrt <code>theta_two</code>. But theta_two is the results of <code>theta_two -= 0.01 * grad</code>, so you get gradients wrt to the result of this operation.<br>
If you want gradients wrt to <code>theta</code>, you should use <code> grad = torch.autograd.grad(loss, theta)[0]</code>. Then you will see that the original value of <code>theta_two</code> is needed for the double backward. And you will need to change to <code>theta_two = theta_two - 0.01 * grad</code>.</p>
<p>Hope this helps.</p>

jwillette: <p>I see now. I wasn’t able to get the graph working, but I got the code snippet working. Below is the final one with the right derivatives for anyone who finds this later. One thing I don’t get though is why<br>
<code>theta -= 0.01 * grad</code> behaves different than <code>theta = theta - 0.01 * grad</code> I though the first one was just a shorthand version and exactly the same. Why did that need to change?</p>
<pre><code class="lang-auto">def maml_simulation() -&gt; None:
    x_i = torch.tensor(3.0, requires_grad=True)
    x_j = torch.tensor(4.0, requires_grad=True)
    theta = torch.tensor(2.0, requires_grad=True)
    loss_func = lambda x: (x - 10) ** 2
    print(f"before training: theta: {theta}")

    prediction = theta * x_i
    loss_one = loss_func(prediction)

    grad = torch.autograd.grad(loss_one, theta, create_graph=False)[0]
    theta_two = theta - 0.01 * grad
    print(
        f"first prediction: {prediction}, loss: {loss_one}, grad: {grad}, theta after update: {theta_two}"
    )

    prediction = theta_two * x_j
    loss = loss_func(prediction)
    grad = torch.autograd.grad(loss, theta)[0]
    print(f"second prediction: {prediction}, loss: {loss}, second grad: {grad}")

</code></pre>

albanD: <p>They are different:</p>
<ul>
<li>The first one modify the Tensor pointed to by <code>theta</code> inplace. So this is now the new value of this Tensor.</li>
<li>
<code>theta = theta - 0.01 * grad</code> creates a new Tensor a associate it with the name “theta”. The Tensor that was originally pointed at by theta is unchanged.</li>
<li>You final code works because you do <code>theta_two = theta - 0.01 * grad</code> and so you keep a reference to the old <code>theta</code> to be able to set it as input for the <code>autograd.grad</code> call.</li>
</ul>