How to add a loss term from hidden layers?

username: DuttaAbhigyan
cooked: <p>Say I want to add a loss term (i.e. gradients from this loss should be propagated) from the output of hidden layers itself i.e. we use the output of a hidden layer and say pass it to a square function (hidden_output/activation)^2. How can I implement this in PyTorch?<br>
Thanks in advance!</p>

username: ptrblck
cooked: <p>I’m not sure if I understand the use case correctly, but you could use any output of a layer and add it to the loss before calling <code>backward</code>.</p>

username: DuttaAbhigyan
cooked: <p>Hi! I’m a new user of PyTorch and I have only used the predefined functions namely, I have first defined outputs from a Neural Net compared it with labels and applied it to criterion nn.CrossEntropyLoss as given in the example for training a simple classifier. So how do we add the outputs from a layer itself?</p>

username: ptrblck
cooked: <p>Here is a small dummy example to use the some intermediate activation in your loss:</p>
<pre><code class="lang-python">class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(10, 64)
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        x1 = F.relu(self.fc1(x))
        x = self.fc2(x1)
        return x, x1


x = torch.randn(10, 10)
y = torch.randn(10, 10)
model = MyModel()
optimizer = optim.SGD(model.parameters(), lr=1e-0)
criterion = nn.MSELoss()

for epoch in range(100):
    optimizer.zero_grad()
    output, aux = model(x)
    loss = criterion(output, y)
    loss = loss + (aux**2).mean()
    loss.backward()
    optimizer.step()
    
    print('Epoch {}, loss {}, aux norm {}'.format(
        epoch, loss.item(), aux.norm()))
</code></pre>
<p>Would this work as a starter for your use case or are you dealing with another problem?</p>

username: DuttaAbhigyan
cooked: <p>Thank You! Very cool trick, I was not aware of it. I think this will do If not I’ll refer you again.</p>

username: Sandeep_Banik
cooked: <p>I believe the aux loss should not impact the weights of fc2 and should only impact fc1. However when I try to compare the weights of 2 networks one with aux and one without aux, the weights are different. I don’t fully understand why did that happen ?</p>
<p>class MyModel(nn.Module):<br>
def <strong>init</strong>(self):<br>
super(MyModel, self).<strong>init</strong>()<br>
self.fc1 = nn.Linear(10, 64)<br>
self.fc2 = nn.Linear(64, 10)</p>
<pre><code>def forward(self, x):
    x1 = F.relu(self.fc1(x))
    x = self.fc2(x1)
    return x, x1
</code></pre>
<p>x = torch.randn(10, 10)<br>
y = torch.randn(10, 10)<br>
model = MyModel()<br>
model2 = MyModel()<br>
model2.load_state_dict(model.state_dict())<br>
optimizer = torch.optim.SGD(model.parameters(), lr=1e-0)<br>
optimizer2 = torch.optim.SGD(model.parameters(), lr=1e-0)<br>
criterion = nn.MSELoss()</p>
<p>for epoch in range(100):<br>
optimizer.zero_grad()<br>
output, aux = model(x)<br>
loss = criterion(output, y)<br>
loss = loss + (aux**2).mean()<br>
loss.backward()<br>
optimizer.step()</p>
<pre><code>optimizer2.zero_grad()
output2, _ = model2(x)
loss2 = criterion(output2, y)
loss2.backward()
optimizer2.step()

print((model2.fc2.weight.detach() == model.fc2.weight.detach()).all())</code></pre>

username: ptrblck
cooked: <p>In your code you are using <code>model.parameters()</code> for both optimizers, so that <code>model2</code> won’t get any updates.<br>
After fixing this, you would expect to see the same parameters after loading the <code>state_dict</code> and after the first weight updates.<br>
Since <code>model.fc1</code> was updated in another way than <code>model2.fc1</code>, you cannot expect the subsequent iterations of the used 100 to yield the same result for <code>.fc2.weight</code>, since the output and thus loss would be different (otherwise there would be no need to use the <code>aux</code> output/loss).</p>
<p>Also, you can post code snippets by wrapping them into three backticks ```, which makes debugging easier. <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>