How to back-propagate gradient through model parameters?

shawn20: <p>I tried to implement a mini-example for <a href="http://proceedings.mlr.press/v70/finn17a/finn17a.pdf" rel="nofollow noopener">MAML</a>, which involves back-propagate through gradients.</p>
<p>Mathematically, I consider a simple linear model ‘output = param * input’ with MSELoss. I would like to do the following update:<br>
’<br>
<strong>param_new</strong> = param - step1 * gradient wrt param of MSELoss(param, train_data)<br>
param = param - step2 *gradient wrt param of MSELoss(<strong>param_new</strong>, test_data)<br>
’<br>
In particular, the term ‘gradient wrt param of MSELoss(<strong>param_new</strong>, test_data)’ is crucial.</p>
<p>I use the following code to compute the desired gradient wrt param of MSELoss(<strong>param_new</strong>, test_data).</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.f = nn.Linear(1,1,bias=False)
        self.init_weight()
    def forward(self, x):
        return self.f(x)
    def init_weight(self):
        nn.init.constant_(self.f.weight, 3.)
m = MyModel()
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(m.parameters(), lr=0.1)

x, t = torch.tensor([2.],requires_grad=True), torch.tensor([2.],requires_grad=True)
x_n, t_n = torch.tensor([3.],requires_grad=True), torch.tensor([3.],requires_grad=True)

y = m(x)
l = loss_fn(y, t)
optimizer.zero_grad()
l.backward(create_graph=True)

for p in m.parameters():
    p_n = p - 0.1*p.grad

y_n = torch.nn.functional.linear(x_n,p_n)
l_n = loss_fn(y_n, t_n)
optimizer.zero_grad()
l_n.backward()
print(m.f.weight.grad)
</code></pre>
<p>I would like to use the model ‘m’ to evaluate ‘y_n’ and ‘l_n’, which corresponds to the following code.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.f = nn.Linear(1,1,bias=False)
        self.init_weight()
    def forward(self, x):
        return self.f(x)
    def init_weight(self):
        nn.init.constant_(self.f.weight, 3.)
m = MyModel()
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(m.parameters(), lr=0.1)

x, t = torch.tensor([2.],requires_grad=True), torch.tensor([2.],requires_grad=True)
x_n, t_n = torch.tensor([3.],requires_grad=True), torch.tensor([3.],requires_grad=True)

y = m(x)
l = loss_fn(y, t)
optimizer.zero_grad()
l.backward(create_graph=True)

state_dict = m.state_dict()
for n, p in m.named_parameters():
    state_dict[n] = p - 0.1*p.grad

m.load_state_dict(state_dict)    
y_n = m(x_n)
l_n = loss_fn(y_n, t_n)
optimizer.zero_grad()
l_n.backward()
print(m.f.weight.grad)
</code></pre>
<p>However, the second snippet cannot compute the desired gradient. From the printed information, it seems that instead of computing ‘gradient wrt param of MSELoss(<strong>param_new</strong>, test_data)’, ‘gradient wrt <strong>param_new</strong> of MSELoss(<strong>param_new</strong>, test_data)’.</p>
<p>I would like to construct a more complicated model (beyond a simple linear one) to do a similar meta update. I was wondering if there is an elegant way to compute the desired gradient which uses the model instead of explicitly calling torch.nn.functional?</p>
<p>Thanks a lot.</p>