Using autograd.grad to take a derivative w.r.t. model parameters

MattTrap: <p>I can’t figure out how to implement this operation with pytorch: <code>\nabla \left( \nabla V_{\theta}^{\top}(s)  \cdot \eta \right)</code></p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/8/d85d2e0509d98b3bc5d9c972e1d646e3e69ee026.png" alt="Screen Shot 2021-06-30 at 2.05.41 PM" data-base62-sha1="uS2LLbWlRmG3AdFDX00rVpwPSpU" width="244" height="58"></p>
<p>V_{\theta} is an ANN that takes in a vector ‘s’, the gradient is the partial w.r.t. weights of the NN and <code>\eta</code> is a the same size as the network weights.</p>
<p>My main question: how to perform the second derivative in pytorch after taking the dot product with ‘\eta’?</p>
<p>I am using autograd.grad(out,in), with in= ANN weights, and the out is the output for some ‘s’. This provides me with the first gradient. But then once I dot the first derivative with ‘\eta’, I am no longer able to computer the gradient with autograd.grad() because there is no computational graph.</p>
<pre><code class="lang-auto"># -*- coding: utf-8 -*-
import torch
import math


dtype = torch.float
device = torch.device("cpu")


# N is batch size; D_in is input dimension;
# H is hidden dimension; D_out is output dimension.
N, D_in, H, D_out = 64, 3, 5, 2

# Create random Tensors to hold inputs and outputs
x = torch.randn(N, D_in, device=device)
y = torch.randn(N, D_out, device=device)

# Use the nn package to define our model as a sequence of layers. nn.Sequential
# is a Module which contains other Modules, and applies them in sequence to
# produce its output. Each Linear Module computes output from input using a
# linear function, and holds internal Tensors for its weight and bias.
model = torch.nn.Sequential(
    torch.nn.Linear(D_in,D_out)
#     torch.nn.Linear(D_in, H),
#     torch.nn.ReLU(),
#     torch.nn.Linear(H, D_out),
)
model.to(device)
# The nn package also contains definitions of popular loss functions; in this
# case we will use Mean Squared Error (MSE) as our loss function.
loss_fn = torch.nn.MSELoss(reduction='sum')

learning_rate = 1e-4
for t in range(100):
    # Forward pass: compute predicted y by passing x to the model. Module objects
    # override the __call__ operator so you can call them like functions. When
    # doing so you pass a Tensor of input data to the Module and it produces
    # a Tensor of output data.
    y_pred = model(x)

    # Compute and print loss. We pass Tensors containing the predicted and true
    # values of y, and the loss function returns a Tensor containing the
    # loss.
    loss = loss_fn(y_pred, y)
    if t % 100 == 99:
        print(t, loss.item())

    # Zero the gradients before running the backward pass.
    model.zero_grad()

    # Backward pass: compute gradient of the loss with respect to all the learnable
    # parameters of the model. Internally, the parameters of each Module are stored
    # in Tensors with requires_grad=True, so this call will compute gradients for
    # all learnable parameters in the model.
    loss.backward()
    print(loss)

    # Update the weights using gradient descent. Each parameter is a Tensor, so
    # we can access its gradients like we did before.
    with torch.no_grad():
        for param in model.parameters():
            param -= learning_rate * param.grad
            
           

outputs = model(torch.tensor([1.,2.5,3.5]))[1]
grad = torch.autograd.grad(outputs, model.parameters(), grad_outputs=None, retain_graph=None, create_graph=False, only_inputs=True, allow_unused=False)
print([x.shape for x in grad])
eta =  [torch.rand(x.shape) for x in grad]
print(eta)
mult = sum([torch.dot(torch.flatten(x),torch.flatten(y)) for x,y in zip(grad,eta)])
mult.requires_grad = True
grad1 = torch.autograd.grad(mult, model.parameters(), grad_outputs=None, retain_graph=None, create_graph=False, only_inputs=True, allow_unused=False)
</code></pre>

soulitzer: <p>You need to have <code>create_graph=True</code> when doing <code>torch.autograd.grad</code> for the first time so that the computation graph is built when computing <code>grad</code>.</p>

MattTrap: <p>I had updated the torch.dot() to requires_grad=True but it gives the error:<br>
<code>One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.</code></p>
<p>I believe this is from the summation. Updated:</p>
<pre><code class="lang-auto">outputs = model(torch.tensor([1.,2.5,3.5]))[1]
grad = torch.autograd.grad(outputs, model.parameters(), grad_outputs=None, retain_graph=True, create_graph=True, only_inputs=True, allow_unused=False)
print(grad)
eta =  [torch.rand(x.shape, requires_grad=True) for x in grad]
print(eta)
mult = torch.sum(torch.tensor([torch.dot(torch.flatten(x),torch.flatten(y)).item() for x,y in zip(grad,eta)], requires_grad=True))
print(mult)
grad1 = torch.autograd.grad(mult, model.parameters(), grad_outputs=None, retain_graph=None, create_graph=True, only_inputs=True, allow_unused=False)
print(grad1)
</code></pre>

soulitzer: <aside class="quote no-group" data-username="MattTrap" data-post="3" data-topic="125494">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/48db29/48.png" class="avatar"> MattTrap:</div>
<blockquote>
<pre><code class="lang-auto">mult = torch.sum(torch.tensor([torch.dot(torch.flatten(x),torch.flatten(y)).item() for x,y in zip(grad,eta)], requires_grad=True))
</code></pre>
</blockquote>
</aside>
<p>You are calling <code>.item</code> and also creating a new tensor on this line both of which break the graph.<br>
You can do something like <code>torch.sum(torch.stack([torch.dot(x.flatten(), y.flatten()) for x, y in zip(grad, eta)]))</code> which should preserve the graph</p>

MattTrap: <p>I had tried with <code>torch.stack()</code> and it didn’t work either. I thought that maybe it is the list comprehension so I tried vectorizing the model resulting grad parameters, but that doesn’t work either. I am not sure if autograd.grad is retaining the graph because when I print it out there is no <code>requires_grad=True</code> attribute:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/b/3b82eb215f136ba91a2a0afde5436a34b29cc455.png" data-download-href="https://discuss.pytorch.org/uploads/default/3b82eb215f136ba91a2a0afde5436a34b29cc455" title="Screen Shot 2021-07-01 at 11.06.10 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3b82eb215f136ba91a2a0afde5436a34b29cc455_2_690x292.png" alt="Screen Shot 2021-07-01 at 11.06.10 AM" data-base62-sha1="8usF3T5sPB7bdfNuwuiLne7Hhad" width="690" height="292" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3b82eb215f136ba91a2a0afde5436a34b29cc455_2_690x292.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3b82eb215f136ba91a2a0afde5436a34b29cc455_2_1035x438.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/b/3b82eb215f136ba91a2a0afde5436a34b29cc455_2_1380x584.png 2x" data-dominant-color="F3E9E9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2021-07-01 at 11.06.10 AM</span><span class="informations">2760×1170 458 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

soulitzer: <p>Whether <code>param</code> will still be in the backward graph of <code>mult</code> will depend on the specifics of the model. For example if we had a linear model, grad of outputs wrt to parameters, is just your model inputs * grad_output, i.e., its not a function of the weight.</p>
<p>But the custom model I defined here is an example of one that does work because the partial derivative of <code>X * k**2</code> wrt k is still <code>2 * X * k</code></p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

model_type = "custom"

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.k = nn.Parameter(torch.rand(1))
    def forward(self, X):
        # the reason we need to clone here is because
        # if an optimizer updates k in-place in no-grad mode it would increment the
        # k's version counter, which could be an issue if k is saved for backward 
        return X * self.k.clone() ** 2

X = torch.rand(10, requires_grad=True).clone()

if model_type == "linear":
    model = nn.Linear(10, 10)
elif model_type == "custom":
    model = Model()

y_hat = model(X).sum()
eta = [torch.rand_like(t, requires_grad=True) for t in model.parameters()]
param_grads = torch.autograd.grad(y_hat, model.parameters(), create_graph=True)
out = torch.max(torch.stack([torch.dot(x.flatten(), y.flatten()) for (x, y) in zip(param_grads, eta)]))

if model_type == "linear":
    # d_out_d_params = torch.autograd.grad(out, model.parameters())
    # RuntimeError: One of the differentiated Tensors appears to not have been used in the graph
    pass
elif model_type == "custom":
    d_out_d_params = torch.autograd.grad(out, model.parameters())
    print(d_out_d_params)
</code></pre>

MattTrap: <p>Thank you for your help! I realized that I needed to have a final non-linear activation function after my last linear layer.<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/2/f/2f01951eeae62f426d35adeeefc286168fa8f5ba.png" alt="Screen Shot 2021-07-08 at 7.33.30 AM" data-base62-sha1="6HPPyHvR1UJLP1b39qdeC2HY3Q6" width="550" height="206"></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/f/cfa5db11ce03f81d8769e5529c691979d67f249a.png" data-download-href="https://discuss.pytorch.org/uploads/default/cfa5db11ce03f81d8769e5529c691979d67f249a" title="Screen Shot 2021-07-08 at 7.32.55 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/f/cfa5db11ce03f81d8769e5529c691979d67f249a_2_690x355.png" alt="Screen Shot 2021-07-08 at 7.32.55 AM" data-base62-sha1="tCWbcueXHRsOaVyt2GHmDKiMUgW" width="690" height="355" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/f/cfa5db11ce03f81d8769e5529c691979d67f249a_2_690x355.png, https://discuss.pytorch.org/uploads/default/optimized/3X/c/f/cfa5db11ce03f81d8769e5529c691979d67f249a_2_1035x532.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/f/cfa5db11ce03f81d8769e5529c691979d67f249a_2_1380x710.png 2x" data-dominant-color="F1F1F1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2021-07-08 at 7.32.55 AM</span><span class="informations">2872×1480 602 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>