Batchnorm in multi-head CNN

username: SofiaCP
cooked: <p>Hey there,</p>
<p>I intend to build a densenet121-based multi-head CNN. My goal is to use the first three denseblocks as a shared set of layers and then create multiple branches using the architecture of the fourth denseblock.</p>
<p>My code is below:</p>
<pre><code class="lang-auto">class Densenet121_v1(nn.Module):
    def __init__(self, num_classes = [1,4,2,8,3]):
        super(Densenet121_v1,self).__init__()
        original_model = torchvision.models.densenet121(pretrained=True)
        self.num_classes=num_classes
        
        self.trunk=original_model.features[:-2]
        
        self.branch1=original_model.features[-2:]
        self.branch2=original_model.features[-2:]
        self.branch3=original_model.features[-2:]
        self.branch4=original_model.features[-2:]
        self.branch5=original_model.features[-2:]
        
        #self.features=nn.Sequential(*self.trunk,*self.branch)
        #self.features = nn.Sequential(*list(original_model.children())[:-1])
        
        self.classifier1=(nn.Linear(1024, self.num_classes[0]))
        self.classifier2=(nn.Linear(1024, self.num_classes[1]))
        self.classifier3=(nn.Linear(1024, self.num_classes[2]))
        self.classifier4=(nn.Linear(1024, self.num_classes[3]))
        self.classifier5=(nn.Linear(1024, self.num_classes[4]))
        
    def forward(self, x):
        
        # shared trunk
        trunk=self.trunk(x)
    
        #block1: 
        f1 = self.branch1(trunk)
        b1 = F.relu(f1, inplace=False)
        b1 = F.adaptive_avg_pool2d(b1, (1, 1)).view(f1.size(0), -1)
        y1 = self.classifier1(b1)
  
        #block2: 
        f2 = self.branch2(trunk)
        b2 = F.relu(f2, inplace=False)
        b2 = F.adaptive_avg_pool2d(b2, (1, 1)).view(f2.size(0), -1)
        y2 = self.classifier2(b2)

        #block3: 
        f3 = self.branch3(trunk)
        b3 = F.relu(f3, inplace=False)
        b3 = F.adaptive_avg_pool2d(b3, (1, 1)).view(f3.size(0), -1)
        y3 = self.classifier3(b3)

        #block4: 
        f4 = self.branch4(trunk) 
        b4 = F.relu(f4, inplace=False)
        b4 = F.adaptive_avg_pool2d(b4, (1, 1)).view(f4.size(0), -1)
        y4 = self.classifier4(b4)

        #block5: 
        f5 = self.branch5(trunk)
        b5 = F.relu(f5, inplace=False)
        b5 = F.adaptive_avg_pool2d(b5, (1, 1)).view(f5.size(0), -1)
        y5 = self.classifier5(b5)
 
        #order predictions
        output= torch.cat((y1,y2[:,0:3],torch.unsqueeze(y4[:,0],1),torch.unsqueeze(y2[:,3],1),
                            y4[:,1:3],torch.unsqueeze(y5[:,0],1),
                            torch.unsqueeze(y3[:,0],1),torch.unsqueeze(y4[:,3],1),
                            torch.unsqueeze(y3[:,1],1),y4[:,4:7],
                            torch.unsqueeze(y5[:,1],1),torch.unsqueeze(y4[:,7],1),
                            torch.unsqueeze(y5[:,2],1)),1)

        
        return output
    
model=Densenet121_v1()
out=model(torch.rand(16,3,224,224))
</code></pre>
<p>When I inspect the computational graph using tensorboard I see that some tensors resulting from batchnorm layers in branch1 are being fed into the remaining branches. This does not happen with the remaining branches:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/f/cfb62badc44493eb15ad3ee58c26b2d3561b2c9b.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/cfb62badc44493eb15ad3ee58c26b2d3561b2c9b" title="final"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/c/f/cfb62badc44493eb15ad3ee58c26b2d3561b2c9b_2_690x313.jpeg" alt="final" data-base62-sha1="tDv8nsRNmmLeD9Mc2s5EXr0AJnl" width="690" height="313" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/c/f/cfb62badc44493eb15ad3ee58c26b2d3561b2c9b_2_690x313.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/c/f/cfb62badc44493eb15ad3ee58c26b2d3561b2c9b_2_1035x469.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/c/f/cfb62badc44493eb15ad3ee58c26b2d3561b2c9b_2_1380x626.jpeg 2x" data-dominant-color="EEEBEB"><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">final</span><span class="informations">1920×872 62.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>How can I avoid this? The heads should be completely independent of each other.</p>

username: InnovArul
cooked: <p>Probably you can try to <code>deepcopy</code> the branches and then resetting the parameters of the <code>branch</code>es to make them independent?<br>
Current code has all the branches sharing the same weights.</p>
<p>Try resetting the parameters as shown here: <a href="https://discuss.pytorch.org/t/how-to-re-set-alll-parameters-in-a-network/20819/12" class="inline-onebox">How to re-set alll parameters in a network - #12 by Brando_Miranda</a></p>

username: SofiaCP
cooked: <p>Thanks for the quick reply.</p>
<p>I’m not sure how to do that. I was trying to do the following:</p>
<pre><code class="lang-auto">def weight_reset(m):
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        m.reset_parameters()

    
class Densenet121_conv(nn.Module):
    def __init__(self, num_classes = [1,4,2,8,3]):
        super(Densenet121_conv,self).__init__()
        original_model = torchvision.models.densenet121(pretrained=True)
        self.num_classes=num_classes
        
        self.trunk=original_model.features[:-2]
        
        self.branch1=original_model.features[-2:]
        self.branch2=original_model.features[-2:]
        self.branch3=original_model.features[-2:]
        self.branch4=original_model.features[-2:]
        self.branch5=original_model.features[-2:]
        
        self.classifier1=(nn.Linear(1024, self.num_classes[0]))
        self.classifier2=(nn.Linear(1024, self.num_classes[1]))
        self.classifier3=(nn.Linear(1024, self.num_classes[2]))
        self.classifier4=(nn.Linear(1024, self.num_classes[3]))
        self.classifier5=(nn.Linear(1024, self.num_classes[4]))
        
    def forward(self, x):
        
        # shared trunk
        trunk=self.trunk(x)
    
        #block1: 
        f1 = self.branch1(trunk)
        b1 = F.relu(f1, inplace=False)
        b1 = F.adaptive_avg_pool2d(b1, (1, 1)).view(f1.size(0), -1)
        y1 = self.classifier1(b1)
        
        # copy and re-set
        f1_copy=copy.deepcopy(f1)
        f1.apply(weight_reset)
  
        #block2: 
        f2 = self.branch2(trunk)
        b2 = F.relu(f2, inplace=False)
        b2 = F.adaptive_avg_pool2d(b2, (1, 1)).view(f2.size(0), -1)
        y2 = self.classifier2(b2)

        #block3: 
        f3 = self.branch3(trunk)
        b3 = F.relu(f3, inplace=False)
        b3 = F.adaptive_avg_pool2d(b3, (1, 1)).view(f3.size(0), -1)
        y3 = self.classifier3(b3)

        #block4: 
        f4 = self.branch4(trunk) 
        b4 = F.relu(f4, inplace=False)
        b4 = F.adaptive_avg_pool2d(b4, (1, 1)).view(f4.size(0), -1)
        y4 = self.classifier4(b4)

        #block5: 
        f5 = self.branch5(trunk)
        b5 = F.relu(f5, inplace=False)
        b5 = F.adaptive_avg_pool2d(b5, (1, 1)).view(f5.size(0), -1)
        y5 = self.classifier5(b5)
 
        output= torch.cat((y1,y2[:,0:3],torch.unsqueeze(y4[:,0],1),torch.unsqueeze(y2[:,3],1),
                            y4[:,1:3],torch.unsqueeze(y5[:,0],1),
                            torch.unsqueeze(y3[:,0],1),torch.unsqueeze(y4[:,3],1),
                            torch.unsqueeze(y3[:,1],1),y4[:,4:7],
                            torch.unsqueeze(y5[:,1],1),torch.unsqueeze(y4[:,7],1),
                            torch.unsqueeze(y5[:,2],1)),1)

        #restore the branch weights after the forward pass
        f1=f1_copy
        
        return output
    
</code></pre>
<p>but got the error <em>Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment</em></p>
<p>Either way, would something like this work during backprop? The idea is that during backprop the common trunk is updated by all losses but the subbranches are only updated by the losses of their outputs.</p>
<p>Further help would be really appreciated.<br>
Thanks!</p>

username: InnovArul
cooked: <aside class="quote no-group" data-username="SofiaCP" data-post="3" data-topic="137258">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sofiacp/48/44302_2.png" class="avatar"> SofiaCP:</div>
<blockquote>
<p>Either way, would something like this work during backprop? The idea is that during backprop the common trunk is updated by all losses but the subbranches are only updated by the losses of their outputs.</p>
</blockquote>
</aside>
<p>Yes this would work. For example, the following work seems to follow this paradigm.<br>
<a href="https://arxiv.org/pdf/1905.09979.pdf">ensemblenet: end-to-end optimization of multi-headed models</a></p>
<aside class="quote no-group" data-username="SofiaCP" data-post="3" data-topic="137258">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sofiacp/48/44302_2.png" class="avatar"> SofiaCP:</div>
<blockquote>
<p>I’m not sure how to do that</p>
</blockquote>
</aside>
<p>This might help.</p><aside class="onebox githubgist" data-onebox-src="https://gist.github.com/InnovArul/871cf866b737b306eb073f3c37a6a8fd">
  <header class="source">

      <a href="https://gist.github.com/InnovArul/871cf866b737b306eb073f3c37a6a8fd" target="_blank" rel="noopener">gist.github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://gist.github.com/InnovArul/871cf866b737b306eb073f3c37a6a8fd" target="_blank" rel="noopener">https://gist.github.com/InnovArul/871cf866b737b306eb073f3c37a6a8fd</a></h4>

  <h5>reset_params.py</h5>
  <pre><code class="Python">import torchvision, copy
import torch, torch.nn as nn

def reset_all_weights(model: nn.Module) -&gt; None:
    """
    refs:
        - https://discuss.pytorch.org/t/how-to-re-set-alll-parameters-in-a-network/20819/6
        - https://stackoverflow.com/questions/63627997/reset-parameters-of-a-neural-network-in-pytorch
        - https://pytorch.org/docs/stable/generated/torch.nn.Module.html
    """</code></pre>
    This file has been truncated. <a href="https://gist.github.com/InnovArul/871cf866b737b306eb073f3c37a6a8fd" target="_blank" rel="noopener">show original</a>

<p>
</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


username: SofiaCP
cooked: <p>Still need to test it with my data set but looking at the graph in tensorboard looks like its working like a charm:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bb546afa8f85fca0d2594049d313da6f1795b276.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/bb546afa8f85fca0d2594049d313da6f1795b276" title="mutantcnn"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bb546afa8f85fca0d2594049d313da6f1795b276_2_690x278.jpeg" alt="mutantcnn" data-base62-sha1="qJc9odaIaCmSSIDisWUFnEQEIey" width="690" height="278" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bb546afa8f85fca0d2594049d313da6f1795b276_2_690x278.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bb546afa8f85fca0d2594049d313da6f1795b276_2_1035x417.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bb546afa8f85fca0d2594049d313da6f1795b276_2_1380x556.jpeg 2x" data-dominant-color="EDECEB"><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">mutantcnn</span><span class="informations">1920×776 56.9 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>Thank you so much!<br>
Sofia</p>

username: SofiaCP
cooked: <p>Hi again,</p>
<p>When you say that without the reset all the branches are sharing the same weights, do you mean that they all start from the same weights? But as training progresses (and backprop happens), they will end up getting different weights, right?</p>
<p>Thank you!</p>