How to calculate the gradient based on another gradient?

lorenmt: <p>Suppose we have two networks with parameters <code>f</code> and <code>g</code> respectively.</p>
<p>And they all have the same input x such that <code>y1, y2=f(x)</code> [multi-task network with two outputs] and <code>z=g(x)</code>.</p>
<p>What I want to do is this: Assume <code>y1=f1(x)</code>, <code>y2=f2(x)</code>, to update <code>f</code> we have two losses in which one is based on the prediction from <code>g(x)</code>, after one gradient update, we want to update <code>g</code> which based on the performance of <code>f1(x)</code>:</p>
<p>f’ = f + \alpha * Grad(L(f1(x), y1_gt) + L(f2(x), g(x)))</p>
<p>g‘ = g + \beta * Grad(L(f1’(x), y1_gt))). # use the updated f’</p>
<p>\alpha, \beta are learning rates.</p>
<p>The current issue is: I don’t know how to implement the <code>grad(L(f1’(x), y1_gt)))</code>.</p>
<p>My current code is (proved to be wrong):</p>
<pre><code># define my model1 with two outputs
class model1(...):
    ...
    def forward(x):
    ...
    return y1, y2

# define my model2 with one output    
class model2(...):
    ...
    def forward(x):
    ...
    return z

def model_fit(pred, gt):
        return [cross_entropy_loss(pred, gt)]

# define optmizers:
optimizer1 = optim.SGD(model1.parameters(), lr=1e-5)
optimizer2 = optim.SGD(model2.parameters(), lr=1e-5)

# run one iteration as a simple example:
y1, y2 = model1(train_data)  # y1,y2=f(x)
z = model2(train_data)   # z=g(x)

optimizer1.zero_grad()
optimizer2.zero_grad()

train_loss1 = model_fit(y1, y1_gt) # compute loss1 L(f1(x), y1_gt)
train_loss2 = model_fit(y2, z) # compute loss2  L(f2(x), g(x))

(train_loss1+train_loss2).backward(create_graph=True) # Grad(L(f1(x), y1_gt) + L(f2(x), g(x))
optimizer1.step() # f’ = f + \alpha * Grad(L(f1(x), y1_gt) + L(f2(x), g(x)))

# 2nd forward pass based on the updated network 1.
y1,y2 = model1(train_data)  # y1,y2 = f'(x)
train_loss1 = model1.model_fit(y1, y1_gt)  # L(f1’(x), y1_gt))
train_loss1.backward()  # Grad(L(f1’(x), y1_gt)))
optimizer2.step()  # g‘ = g + \beta * Grad(L(f1’(x), y1_gt)))
</code></pre>
<p>However, the performance didn’t go as expected. It only calculates the gradient on train_loss2 in the first step() and retains the same gradient on the second step().</p>
<p>Update: I think I have understood what my code is doing: in the first backward function: it calculates the gradient w.r.t. both model1 and model2, the optimizer1.step() updates the gradients in model1 <strong>(Grad(L(f1(x), y1_gt) + L(f2(x), g(x))) w.r.t. f)</strong>. However, in the second forward pass, the model2 recreates the computational graph and it cannot retain the previous graph. Thus, the gradient of train_loss1 w.r.t. model2 would be simply zeros. Since I put a create_graph=True in the first gradient update, the second backward pass will cumulate the gradients, thus the same as train_loss2 gradient w.r.t. model2 in the first gradient update <strong>(Grad(L(f2(x), g(x)) w.r.t. g, (accumulates grad(L(f1(x), y1_gt) w.r.t. g which is zero) )</strong>.</p>
<p>Now the problem becomes: how to retain the graph in the second forward pass so that we can compute the higher order gradients.</p>
<p>Any comments are welcomed. Thanks!</p>

WeiQin-C: <p>Hi, <a class="mention" href="/u/lorenmt">@lorenmt</a>  have you solved this problem?</p>