How to compute 2nd derivative of the loss computing from a fine-tuned model w.r.t pre-trained model's input?

username: hkz
cooked: <p>I’m new to Pytorch. I have a pre-trained model W1 and two data point (x1, y1), (x2, y2). Firstly, I fine-tune W1 with (x1, y1) in one epoch, thus the model update <img src="https://discuss.pytorch.org/uploads/default/original/3X/7/7/771fe9804c97881f24d1ad69d0b1c81b0105018d.png" alt="mylatex20191228_223045" data-base62-sha1="gZPbNN215pQjEjAHEtRkZmVlxHT" width="47" height="20">   can be obtained by doing<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/6/f/6f1f978c3cd6e99c2de916832b40867e37e745a4.png" alt="mylatex20191228_223001" data-base62-sha1="fR2Gs6WnvKX3VUBQNSpRnmZsgks" width="260" height="52"><br>
And then, I get the test loss for (x2, y2) w.r.t fine-tuned model <img src="https://discuss.pytorch.org/uploads/default/original/3X/a/2/a29f5e5a82cfb7a9d7dd43c65c19e4a7f7320f86.png" alt="mylatex20191228_223142" data-base62-sha1="ncCNl2whe7X9Gv6899ORYFVzBBA" width="168" height="20">  by<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/a/3/a35ba8298759e1ced1b7da62d12c69f1ac05a16f.png" alt="mylatex20191228_215843" data-base62-sha1="nj8cv926h5NAiozrsEzVp5JzZaT" width="298" height="25"><br>
What I want to do is to use the test_loss to update x1 by doing<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bb47b40c6c79204038b44ba4e2ba34283bf389c3.png" data-download-href="https://discuss.pytorch.org/uploads/default/bb47b40c6c79204038b44ba4e2ba34283bf389c3" title="mylatex20191228_222925"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bb47b40c6c79204038b44ba4e2ba34283bf389c3.png" alt="mylatex20191228_222925" data-base62-sha1="qIKUspvOX6r9JgfdQooIm5xXz0v" width="690" height="51" data-dominant-color="000000"><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">mylatex20191228_222925</span><span class="informations">691×52 2.44 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>
<p>Here I offer a toy code where the model is F(x)=wx, and the two data points are (1, 1) and (2, 2) respectively.</p>
<pre><code class="lang-auto">import torch, torch.nn as nn
from torch.autograd import grad

torch.manual_seed(1)
torch.cuda.manual_seed(1)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc_blocks = nn.Sequential(
            nn.Linear(1, 1, bias=False)
        )
    def forward(self, x):
        output = self.fc_blocks(x)
        return output

def weight_init(m):
    if isinstance(m, nn.Linear):
        nn.init.constant_(m.weight, 2.0)

net = Net()
net.apply(weight_init)
criterion = nn.MSELoss()

x_1 = torch.tensor([[1.0]], requires_grad=True)
y_1 = torch.tensor([[1.0]])
x_2 = torch.tensor([[2.0]])
y_2 = torch.tensor([[2.0]])

x_optimizer = torch.optim.SGD([x_1,], lr = 1)
net_optimizer = torch.optim.SGD(net.parameters(), lr = 1)

for i in range(1):
    def closure():
        net_optimizer.zero_grad()
        y_pred = net(x_1)
        pred_loss = criterion(y_pred, y_1)
        dp_dw = grad(pred_loss, net.parameters(), create_graph=True, retain_graph=True)
        dp_dw[0].backward()
        print(x_1.grad)
        
        for param in net.parameters():
            param.data = param.data - param.grad.data

        y_test = net(x_2)
        test_loss = criterion(y_test, y_2)
        test_loss.backward(create_graph=True)
        print(x_1.grad)
        return test_loss
    x_optimizer.step(closure)
</code></pre>
<p>And the result is</p>
<pre><code class="lang-auto">$ python test.py
tensor([[6.]], grad_fn=&lt;CloneBackward&gt;)
tensor([[6.]], grad_fn=&lt;CloneBackward&gt;)
</code></pre>
<p>It seems that the x1.grad doesn’t change in step2 mainly caused by the the computation graph forgetting this operation <img src="https://discuss.pytorch.org/uploads/default/original/3X/a/2/a29f5e5a82cfb7a9d7dd43c65c19e4a7f7320f86.png" alt="mylatex20191228_223142" data-base62-sha1="ncCNl2whe7X9Gv6899ORYFVzBBA" width="168" height="20">. But I have manually re-assign the model’s parameters intead of using in-place operation implemented in <code>step()</code> function.<br>
What am I doing wrong here? Any help is appreciated. Thank you.</p>

username: albanD
cooked: <p>Hi,</p>
<p>The main problem you encounter here is that the <code>nn.Module</code> and <code>torch.optim</code> modules are not built to do what you want. In particular, the parameters of the modules and the optimizers are built such that weight updates are not differentiated through.</p>
<p>The two issues you want to fix here are:</p>
<ul>
<li>You don’t want to use builtin optimizers as they perform their updates in a non-differentiable way.</li>
<li>You don’t want to use vanilla modules because their weights are <code>nn.Parameter</code> for which you never track history (they are always leaf Tensors).</li>
</ul>
<p>If you only want to perform testing for research purposes, I would advise</p>
<ul>
<li>Create regular pytorch modules.</li>
<li>Create the new base parameters as <code>mod.true_weight = nn.Parameter(mod.weight)</code>
</li>
<li>Delete all the old parameters by doing <code>del mod.weight</code> manually.</li>
<li>Have a way to manually set the <code>.weight</code> to the base parameters (that will use only for your first iteration)
<ul>
<li>For every param, do <code>mod.weight = mod.true_weight.view_as(mod.true_weight)</code> (do an operation to make sure you don’t get a new Tensor (that share memory)).</li>
</ul>
</li>
<li>Have a custom optimizers that will do <em>in a differentiable way</em> the weight update based on the true_weight’s gradients: <code>mod.weight = mod.weight + lr * mod.true_weight.grad</code>.</li>
<li>Use the regular <code>mod(inp)</code> to do the forward calls.</li>
</ul>
<p>That way, you will be able to manually do these different iterations.</p>

username: hkz
cooked: <p>Thanks for your patient reply and please forgive my slow response.</p>
<p>I know the idea in your operations, but I don’t know how to implement it. Could you please provide me some specific examples?</p>
<p>By the way, I wonder if it’s possible to manually compute the gradient<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/7/e7e52005eb17dac0a3f40e90cda801ea6d505273.png" data-download-href="https://discuss.pytorch.org/uploads/default/e7e52005eb17dac0a3f40e90cda801ea6d505273" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/7/e7e52005eb17dac0a3f40e90cda801ea6d505273_2_690x69.png" alt="image" data-base62-sha1="x5rc1pxKRCKhVsFJ4qPa7alX25Z" width="690" height="69" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/7/e7e52005eb17dac0a3f40e90cda801ea6d505273_2_690x69.png, https://discuss.pytorch.org/uploads/default/original/3X/e/7/e7e52005eb17dac0a3f40e90cda801ea6d505273.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/7/e7e52005eb17dac0a3f40e90cda801ea6d505273.png 2x" data-dominant-color="EFEFEF"><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">image</span><span class="informations">900×91 10.8 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><br>
By <code>dp_dw[0].backward()</code> and <code>x_1.grad</code> in the first iteration, I can get the gradient <img src="https://discuss.pytorch.org/uploads/default/original/3X/b/0/b051c7a1c392e9967d1c8feb7570204be071b42e.png" alt="image" data-base62-sha1="p9NfknOAuWzzH5b0LAQiqpsd6KW" width="165" height="90"> (Is it correct since \Delta W1 is a function of x1 while W1 is not a function of x? ) And it’s easy to get the gradient <img src="https://discuss.pytorch.org/uploads/default/original/3X/d/3/d396d45da5fb90b746dae45aac97fb76736a0926.png" alt="image" data-base62-sha1="ubNTCqyKFVOUBIqOvoaD9a1eWl8" width="328" height="92"> with <code>test_loss.backward()</code> and <code>net.parameters().grad</code> in the second iteration. So can I pass the manually computed gradient to a optimizer in order to optimize x1?</p>