Calculating Hessian vectors for outputs with respect to the inputs

keunwoo: <p>Hi. I want to calculate the second derivatives of outputs w.r.t. inputs.<br>
I found some codes that calculate Hessian matrices, such as <a class="mention" href="/u/apaszke">@apaszke</a> 's <a href="https://gist.github.com/apaszke/226abdf867c4e9d6698bd198f3b45fb7" rel="nofollow noopener">code</a>.</p>
<p>As I understand, <a class="mention" href="/u/apaszke">@apaszke</a> 's code calculates Hessian matrix for all elements in <code>y</code> w.r.t. all elements in <code>x</code>. But what I want is the second derivative of <code>y</code> w.r.t. to the corresponding input <code>x</code>. So I changed some of their code.</p>
<p>And I tested that code as below.</p>
<pre><code class="lang-auto">import torch.nn as nn
import torch

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = x ** 3
        return x

class SigmoidNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, 1), nn.Sigmoid())

    def forward(self, x):
        return self.net(x)

class ReLUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, 1), nn.ReLU())

    def forward(self, x):
        return self.net(x)

def jacobian_vector(y, x, create_graph = False):
    """
        reference: https://gist.github.com/apaszke/226abdf867c4e9d6698bd198f3b45fb7
    """
    jac = []
    batch_size = y.size(0)
    flat_y = y.reshape(-1)
    grad_y = torch.zeros_like(flat_y)
    for i in range(batch_size):
        grad_y[i] = 1.
        grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph)
        grad_x = grad_x.reshape(x.shape)
        jac_vec = grad_x[i] # y[i] = f(x[i]). so only grad_x[i] is valid for y[i]
        jac.append(jac_vec)
        grad_y[i] = 0.
    return torch.stack(jac, axis = 0)

def hessian(y, x):
    first_dev = jacobian_vector(y, x, create_graph=True)
    print("first dev:\n", first_dev)
    return jacobian_vector(first_dev, x)

if __name__=="__main__":
    print("Test Hessian Function")
    torch.autograd.set_detect_anomaly(True)

    x = torch.Tensor([[1,1], [3,3], [2,2]]).requires_grad_(True)
    net = SimpleNet()
    # net = SigmoidNet() # bug

    y = net(x)

    print("x:\n",x)
    print("y:\n",y)

    h = hessian(y, x)

    print("hessian:\n",h)

</code></pre>
<p>I have two questions about this code.</p>
<ol>
<li>When I use <code>SigmoidNet()</code>, the code shows an error that I modified a variable in place. But I cannot find where this modification is happening.</li>
</ol>
<p>This is the error message with <code>torch.autograd.set_detect_anomaly(True)</code></p>
<pre><code class="lang-auto">[W python_anomaly_mode.cpp:60] Warning: Error detected in SigmoidBackwardBackward. Traceback of forward call that caused the error:
  File "hessian_test.py", line 63, in &lt;module&gt;
    h = hessian(y, x)
  File "hessian_test.py", line 46, in hessian
    first_dev = jacobian_vector(y, x, create_graph=True)
  File "hessian_test.py", line 38, in jacobian_vector
    grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph)
  File "(deleted)/python3.6/site-packages/torch/autograd/__init__.py", line 192, in grad
    inputs, allow_unused)
 (function print_stack)
Traceback (most recent call last):
  File "hessian_test.py", line 63, in &lt;module&gt;
    h = hessian(y, x)
  File "hessian_test.py", line 48, in hessian
    return jacobian_vector(first_dev, x)
  File "hessian_test.py", line 38, in jacobian_vector
    grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph)
  File "(deleted)/python3.6/site-packages/torch/autograd/__init__.py", line 192, in grad
    inputs, allow_unused)
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [3, 1]] is at version 6; expected version 5 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!
</code></pre>
<ol start="2">
<li>Its result is weird. When I use <code>SimpleNet()</code> , the expected answer is</li>
</ol>
<pre><code class="lang-auto">first dev: # this should be 3x**2
 tensor([[3., 3.],
        [27., 27.],
        [12., 12.]], grad_fn=&lt;StackBackward&gt;)
hessian: # this should be 6x
 tensor([[6., 6.],
        [18., 18.],
        [12., 12.]])
</code></pre>
<p>However, the actual output is</p>
<pre><code class="lang-auto">Test Hessian Function
x:
 tensor([[1., 1.],
        [3., 3.],
        [2., 2.]], requires_grad=True)
y:
 tensor([[ 1.,  1.],
        [27., 27.],
        [ 8.,  8.]], grad_fn=&lt;PowBackward0&gt;)
first dev:
 tensor([[3., 0.],
        [0., 0.],
        [0., 0.]], grad_fn=&lt;StackBackward&gt;)
hessian:
 tensor([[6., 0.],
        [0., 0.],
        [0., 0.]])
</code></pre>
<p>I guess I don’t have a clear understanding of <code>torch.autograd.grad</code>.<br>
What am I doing wrong?<br>
Moreover, I don’t understand what is <code>grad_output</code> of <code>autograd.grad</code>. The <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.grad" rel="nofollow noopener">documentation</a> isn’t clear to me.</p>

keunwoo: <p>I guess I found the solution.<br>
What I wanted to do was to calculate <a href="https://en.wikipedia.org/wiki/Laplace%27s_equation" rel="nofollow noopener">Laplace’s equation</a>.</p>
<ol>
<li>
<p>I found <a href="https://discuss.pytorch.org/t/hessian-of-output-with-respect-to-inputs/60668/10">this trick</a> to solve the error.</p>
</li>
<li>
<p>I was misunderstanding <a class="mention" href="/u/apaszke">@apaszke</a> 's <a href="https://gist.github.com/apaszke/226abdf867c4e9d6698bd198f3b45fb7" rel="nofollow noopener">code</a>.</p>
</li>
</ol>
<p>I fixed my code and it seems like what I want. The code is clumsy, but at least it works.</p>
<pre><code class="lang-auto">import torch.nn as nn
import torch

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = x ** 3
        return x

class SimpleSumNet(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        x = x ** 3
        return x.sum(dim=1).unsqueeze(1)

class SigmoidNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, 1), nn.Sigmoid())

    def forward(self, x):
        x = self.net(x)
        return x
class ReLUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, 1), nn.ReLU())

    def forward(self, x):
        return self.net(x)

def jacobian(y, x, create_graph = False):
    """
        reference: https://gist.github.com/apaszke/226abdf867c4e9d6698bd198f3b45fb7
    """
    jac = []
    flat_y = y.reshape(-1)
    grad_y = torch.zeros_like(flat_y)
    grad_y = torch.zeros_like(flat_y)
    for i in range(len(flat_y)):
        grad_y[i] = 1.
        grad_x, = torch.autograd.grad(flat_y, x, grad_y, retain_graph=True, create_graph=create_graph)
        grad_x = grad_x.reshape(x.shape)
        jac.append(grad_x.reshape(x.shape))
        grad_y[i] = 0.

    return torch.stack(jac, axis = 0).reshape(y.shape + x.shape)

def laplace(y, x):
    """ Laplace's Equation
        https://en.wikipedia.org/wiki/Laplace%27s_equation
    """
    jac1 = jacobian(y, x, create_graph=True)

    batch_size = jac1.size(0)
    jvs = []
    for i in range(batch_size):
        jvs.append(jac1[i, :, i, :])

    first_dev = torch.stack(jvs, dim=0)

    jac2 = jacobian(first_dev, x)

    jvs = []
    for i in range(batch_size):
        jvs.append(jac2[i, :, :,  i, :])

    second_dev = torch.stack(jvs, dim=0)
    return second_dev.sum(dim = -1).sum(dim = -1)

if __name__=="__main__":

    x = torch.Tensor([[1,1], [3,3], [2,2]]).requires_grad_(True)
    net = SigmoidNet()

    y = net(x)**1 # trick: https://discuss.pytorch.org/t/hessian-of-output-with-respect-to-inputs/60668/10

    print("x:\n",x)
    print("y:\n",y)

    l = laplace(y, x)
    print("laplace:\n", l, l.shape)

</code></pre>
<p>I would like to know a more elegant way if there is any.</p>