Distilation using hooks

azso: <p>Hello All,<br>
I am trying to apply distilation to my model. For this question purpose, I have prepared toy example, and was wondering, why it behaves on a different ways:</p>
<ol>
<li>first example i am getting additional hidden output explicitly</li>
<li>second example i am getting additional hidden output by setting hook-function there.</li>
</ol>
<p>My expectation is that both of them should give the same results.</p>
<p><strong>The question is:</strong><br>
Please could someone explain, whats wrong (it it is) with hook way for getting hidden output?</p>
<pre><code class="lang-auto">    # common inputs and initial data
    import torch
    import torch.nn.functional as F
    from torch import nn, optim
    from matplotlib import pyplot as plt
    
    x = torch.randn(10, 10)
    y = torch.randn(10, 10)
</code></pre>
<p><strong>First example:</strong></p>
<pre><code class="lang-auto">   # getting hidden output explicitly
    class MyModel(nn.Module):
        def __init__(self):
            super(MyModel, self).__init__()
            self.fc1 = nn.Linear(10, 64)   # hidden layer for getting output
            self.fc2 = nn.Linear(64, 10)
    
        def forward(self, x):
            x1 = F.relu(self.fc1(x))
            x = self.fc2(x1)
            return x, x1
    
    model = MyModel()
    optimizer = optim.SGD(model.parameters(), lr=1e-0)
    criterion = nn.MSELoss()
    losses = []
    
    for epoch in range(100):
        optimizer.zero_grad()
        output, aux = model(x)
        loss = criterion(output, y)
        loss = loss + (aux**2).mean()
        losses.append(loss.item())
        loss.backward()
        optimizer.step()
</code></pre>
<p><strong>Second example:</strong></p>
<pre><code class="lang-auto">    class MyDistilModel(nn.Module):
        def __init__(self):
            super(MyDistilModel, self).__init__()
            self.fc1 = nn.Linear(10, 64)     # the same hidden layer for getting output
            self.fc2 = nn.Linear(64, 10)
    
        def forward(self, x):
            x1 = F.relu(self.fc1(x))
            x = self.fc2(x1)
            return x
    
    d_model = MyDistilModel()
    d_optimizer = optim.SGD(d_model.parameters(), lr=1e-0)
    criterion = nn.MSELoss()
    
    hidden = None    # will collecting hidden output here
    
    def hook(module, input_, output):
        global hidden 
        hidden = output
        
    lyrs = []
    for l in d_model.children():
        lyrs.append(l)
        print(l)
        
    lyrs[0].register_forward_hook(hook)  # setting hook to the very first layer 
    
    d_losses = []
    for epoch in range(100):
        d_optimizer.zero_grad()
        output = d_model(x)
        loss = criterion(output, y)
        loss = loss + (hidden**2).mean()
        loss.backward()
        d_optimizer.step()
        d_losses.append(loss.item())
</code></pre>
<p>Here the visualization of <code>losses</code> and <code>d_losses</code> difference</p>
<pre><code class="lang-auto">    f, (ax1, ax2, ax3) = plt.subplots(3,1)
    f.set_figheight(10)
    f.set_figwidth(7)
    ax1.plot(list(range(100))[20:], losses[20:], label = 'losses')
    ax2.plot(list(range(100))[20:], d_losses[20:], label = 'd_losses')
    ax3.plot(list(range(100))[20:], losses[20:], label = 'losses')
    ax3.plot(list(range(100))[20:], d_losses[20:], label = 'd_losses')
    ax1.legend()
    ax2.legend()
    ax3.legend()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/a/eae7d93ee40a28fd227ac36a757ecf0894d986e0.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/eae7d93ee40a28fd227ac36a757ecf0894d986e0" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/eae7d93ee40a28fd227ac36a757ecf0894d986e0_2_401x499.jpeg" alt="image" data-base62-sha1="xw4sBPl1jk903EXdsaKBqQr2SRi" width="401" height="499" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/eae7d93ee40a28fd227ac36a757ecf0894d986e0_2_401x499.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/eae7d93ee40a28fd227ac36a757ecf0894d986e0_2_601x748.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/a/eae7d93ee40a28fd227ac36a757ecf0894d986e0_2_802x998.jpeg 2x" data-dominant-color="F9F9FA"><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">994×1238 89.5 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><strong>Upd.</strong><br>
Actually i have fixed it. The difference between two examples was in a point for getting hidden output. So in first example, <code>F.ReLU(x)</code> was additional applied to getting <code>x1</code>. So examples are not the same. After fixed it, it gives the same results</p>