How to compute the gradient of gradient if I have two models?

username: WeiQin-C
cooked: <p>Hi, I am working on a problem where I have two models, namely a Teacher model (A) and a student model (B).</p>
<h1>Phase 1</h1>
<p>The Teacher network is used to generate pseudo-labels for a set of unlabelled train set <code>X1</code>. The pseudo-labels are used as ground truth to train the student network. The student network is updated based on the loss computed using the prediction from student network and the pseudo-labels.</p>
<h1>Phase 2</h1>
<p>Given labelled train set, <code>(X2, Y2)</code>, we use the updated student model B2 and perform forward propagation. Then a loss is computed between the prediction from B2 and Y2. Rather than updating the student network again, now I would like to update the Teacher model A. In other words, I would like to compute the gradient of loss in phase 2 w.r.t model A. However, in my implementation, when I do loss.backward() in phase 2, only the gradient w.r.t model B is computed. May I ask how I can compute the gradient w.r.t model A???</p>
<p>Below is a much more detailed explanation and I have also pasted my code at the very end.</p>
<p>Given a set of inputs, X1,</p>
<pre><code class="lang-auto">Y_A1 = A(X1)
Y_B1 = B(X1)
loss1 = crossEntropy(Y_A1, Y_B1)
loss1.backward()
</code></pre>
<p>By calling loss1.backward() will computes the gradient of <code>loss1</code> w.r.t all the parameters in A and B. Then to update only model B,</p>
<pre><code class="lang-auto">optimB.step() 
B -&gt; B2
</code></pre>
<p>is called, where optimB can be something like</p>
<pre><code class="lang-auto">optimB = torch.optim.SGD(B.parameters(), lr=0.01)
</code></pre>
<p>Then, given another set of input X2, and labels Y2</p>
<pre><code class="lang-auto">Y_B2 = B2(X2)
loss2 = crossEntropy(Y_B2, Y2)
loss2.backward()
</code></pre>
<p>I was under the assumption that by calling <code>loss2.backward()</code>, gradient w.r.t. to parameters in model A will also be computed (second order derivative of <code>loss1</code> w.r.t model A) because</p>
<pre><code class="lang-auto">Y_B2 = B2(X2)
B2 = B - lr * dloss1/dB

To update A -&gt; A2
A2 = A - lr * dloss2/dA
A2 = A - lr * d(crossEntropy(Y_B2, Y2))/dA
A2 = A - lr * d(crossEntropy(B2(X2), Y2))/dA
</code></pre>
<p>Therefore, in order to update A we need to compute <code>d(dloss1/db)/dA</code>. However, in my implementation, when I called <code>loss2.backward()</code>, only gradient of <code>dloss2/dB</code> is computed.</p>
<p>I have pasted my code below:</p>
<pre><code class="lang-auto">X1 = torch.randn(5,2, dtype=torch.float32)
X2 = torch.randn(5,2, dtype=torch.float32)
Y2 = torch.randint(0,2, (5,1), dtype=torch.float32)

# define linear model
class linearModel(nn.Module):
    def __init__(self):
        super(linearModel, self).__init__()

        self.layer1 = nn.Linear(2, 2, bias=False)
        self.layer2 = nn.Linear(2, 1, bias=False)
        self.activation = nn.Tanh()

    def forward(self, x):
        x = self.layer1(x)
        x = self.activation(x)
        x = self.layer2(x)
        return torch.sigmoid(x)

# define loss function: Cross Entropy
def loss_fcn(pred, target):
     loss = -torch.mean(target*torch.log(pred) - (1-target)*torch.log(1-pred))
  return loss

modelA = linearModel()
modelB = linearModel() 

optimA = torch.optim.SGD(modelA.parameters(), lr=0.01)
optimB = torch.optim.SGD(modelB.parameters(), lr=0.01)

# Phase 1
Y_A1 = modelA(X1)
Y_B1 = modelB(X1)
loss1 = loss_fcn(Y_B1, Y_A1)

# Compute gradient w.r.t. B and update modelB 
loss1.backward()
optimB.step()

# Phase 2
Y_B2 = modelB(X2)   # model B has been updated 
loss2 = loss_fcn(Y_B2, Y2)
# Compute gradients w.r.t. model A and update model A
loss2.backward()
optimA.step()
</code></pre>
<p>However, <code>modelA.layer1.weight.grad</code> returns <code>None</code>.</p>

username: zcajiayin
cooked: <p>Based on your code, your <code>modelA()</code>'s parameters didn’t get involved in the computation graph of calculating <code>loss2</code>, so when you print it gives you <code>None</code></p>

username: WeiQin-C
cooked: <p>Hi, thanks for your reply. Is there a way to manually compute the gradients of loss2 w.r.t modelA’s parameters? Thank you</p>

username: zcajiayin
cooked: <p>In this case, you should pass <code>X2</code> to <code>modelA()</code>:</p>
<pre><code class="lang-auto">Y_A2 = modelA(X2)
loss2 = loss_fnc(Y_A2, Y2)
loss2.backward()
optimA.step()
</code></pre>