How to implement two networks with shared weights and with separate Batch Normalizations?

f10w: <p>Hello,</p>
<p>I am working on a similar network similar to shared weight networks:</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/2X/d/de49335b758b479c6cebbb6dbcaf0389f5925928.jpg" alt="shared_network" data-base62-sha1="vIqPsDZIcLqFtlU3uyZiCVX48cM" width="525" height="327"></p>
<p>This image is obtained from the following blog post, in which the author mentioned implementing <strong>separate batch normalizations</strong> for each network:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://www.alexirpan.com/2017/04/26/perils-batch-norm.html">
  <header class="source">
      <img src="https://www.alexirpan.com/favicon-32x32.png" class="site-icon" width="" height="">

      <a href="https://www.alexirpan.com/2017/04/26/perils-batch-norm.html" target="_blank" rel="noopener nofollow ugc">alexirpan.com</a>
  </header>

  <article class="onebox-body">
    

<h3><a href="https://www.alexirpan.com/2017/04/26/perils-batch-norm.html" target="_blank" rel="noopener nofollow ugc">On The Perils of Batch Norm</a></h3>

  <p>This post is written for deep learning practitioners, and assumes you
know what batch norm is and how it works.</p>


  </article>

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

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

<p>Any ideas on how to implement that in PyTorch please?<br>
Thank you very much in advance for your help!</p>

allenye0119: <p>Maybe you can pass an additional input to the <code>forward</code> function indicating the image type and decide which batch norm to use, a simple example code:</p>
<pre><code class="lang-auto">import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(...)
        self.conv2 = nn.Conv2d(...)
        self.conv3 = nn.Conv2d(...)
        self.bn1 = nn.BatchNorm2d(...)
        self.bn2 = nn.BatchNorm2d(...)
   
    def forward(x, data_type):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        if data_type == 'type1':
            x = self.bn1(x)
        elif data_type == 'type2':
            x = self.bn2(x)

        return x
</code></pre>

f10w: <p>Thanks, <a class="mention" href="/u/allenye0119">@allenye0119</a>!<br>
I can see the idea now, but I don’t see how to use the modified <code>forward()</code> function. The training code at each epoch looks like this:</p>
<pre><code>    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
</code></pre>
<p>So I guess I have to modify <code>Model()</code> so that it takes an additional argument: <code>output = model(data, data_type)</code>? Thanks again for your help.</p>

allenye0119: <p>You don’t have to modify anything other than the <code>forward</code> function.<br>
Just do <code>output = model(data, data_type)</code>.</p>

f10w: <p>Oh, cool!<br>
Great help. Thank you very much, <a class="mention" href="/u/allenye0119">@allenye0119</a>! <img src="//discuss.pytorch.org/images/emoji/apple/smiley.png?v=5" title=":smiley:" class="emoji" alt=":smiley:"></p>

f10w: <p>Hi <a class="mention" href="/u/allenye0119">@allenye0119</a>. I tried your idea on a more complex model and it is not easy at all <img src="//discuss.pytorch.org/images/emoji/apple/frowning.png?v=5" title=":frowning:" class="emoji" alt=":frowning:">  For example the following ResNet18:</p>
<pre><code>class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(in_planes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)

        self.shortcut = nn.Sequential()
        if stride != 1 or in_planes != self.expansion*planes:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(self.expansion*planes)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out

class ResNet(nn.Module):
    def __init__(self, block, num_blocks, num_classes=10):
        super(ResNet, self).__init__()
        self.in_planes = 64

        self.conv1 = conv3x3(3,64)
        self.bn1 = nn.BatchNorm2d(64)
        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)
        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=2)
        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)
        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)
        self.linear = nn.Linear(512*block.expansion, num_classes)

    def _make_layer(self, block, planes, num_blocks, stride):
        strides = [stride] + [1]*(num_blocks-1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_planes, planes, stride))
            self.in_planes = planes * block.expansion
        return nn.Sequential(*layers)

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = F.avg_pool2d(out, 4)
        out = out.view(out.size(0), -1)
        out = self.linear(out)
        return out


def ResNet18():
    return ResNet(BasicBlock, [2,2,2,2])
</code></pre>
<p>I guess I have to be able to do <code>self.shortcut(x, data_type)</code> and <code>self.layer1(out, data_type)</code> but I have no idea how to do that. I can unroll the blocks and the layers to get the 18 layers explicitly and then modify the forward function for each layer, but I don’t think it is a reasonable solution.</p>
<p>Do you have any ideas on that? Thank you so much again!</p>

allenye0119: <p>I think the main difficulty stems from the use of <code>nn.Sequential</code>, which only takes in a single argument <code>input</code>. So it’s not possible to do branching in <code>nn.Sequential</code> when additional info is required (in this case, the data type). Which means that we need to replace all <code>nn.Sequential</code>.</p>
<p>For <code>BasicBlock</code>:</p>
<pre><code class="lang-auto">class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_planes, planes, stride=1):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(in_planes, planes, stride)
        self.bn1 = nn.BatchNorm2d(planes)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = nn.BatchNorm2d(planes)

        # self.shortcut = nn.Sequential()
        # if stride != 1 or in_planes != self.expansion*planes:
        #     self.shortcut = nn.Sequential(
        #         nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False),
        #         nn.BatchNorm2d(self.expansion*planes)
        #     )
        if stride != 1 or in_planes != self.expansion*planes:
            self.identity_shortcut = False
            self.shortcut_conv = nn.Conv2d(in_planes, self.expansion*planes, kernel_size=1, stride=stride, bias=False)
            self.shortcut_bn1 = nn.BatchNorm2d(self.expansion*planes)
            self.shortcut_bn2 = nn.BatchNorm2d(self.expansion*planes)
        else:
            self.identity_shortcut = True

    def forward(self, x, data_type):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))

        # out += self.shortcut(x)
        if not self.identity_shortcut:
            x = self.shortcut_conv(x)
            if data_type == 'type1':
                x = self.sortcut_bn1(x)
            elif data_type == 'type2':
                x = self.sortcut_bn2(x)
        out += x

        out = F.relu(out)
        return out
</code></pre>
<p>For <code>ResNet._make_layer</code>:</p>
<pre><code class="lang-auto">def _make_layer(self, block, planes, num_blocks, stride):
    strides = [stride] + [1]*(num_blocks-1)
    layers = []
    for stride in strides:
        layers.append(block(self.in_planes, planes, stride))
        self.in_planes = planes * block.expansion
    return nn.ModuleList(layers)
</code></pre>
<p>and <code>ResNet.forward</code>:</p>
<pre><code class="lang-auto">def forward(self, x, data_type):
    out = F.relu(self.bn1(self.conv1(x)))
    # out = self.layer1(out)
    # out = self.layer2(out)
    # out = self.layer3(out)
    # out = self.layer4(out)
    for layer in self.layer1:
        out = layer(out, data_type)
    for layer in self.layer2:
        out = layer(out, data_type)
    for layer in self.layer3:
        out = layer(out, data_type)
    for layer in self.layer4:
        out = layer(out, data_type)
    out = F.avg_pool2d(out, 4)
    out = out.view(out.size(0), -1)
    out = self.linear(out)
    return out
</code></pre>
<p>This is only a quick fix, I’m sure the code can be made more elegant.</p>

f10w: <p>You’re awesome, <a class="mention" href="/u/allenye0119">@allenye0119</a>! <img src="//discuss.pytorch.org/images/emoji/apple/smiley.png?v=5" title=":smiley:" class="emoji" alt=":smiley:"><br>
That’s already much more elegant than I expected! Thanks again!</p>

tau: <p>Hi, I also want to implement a split batch norm layer and I can decide to use which in different forward steps. I want to do something like these:</p>
<pre><code class="lang-auto">output1=model(data1)
model.switch("bn2")
output2=model(data2)
model.switch("bn1")
output3=model(data3)
...
</code></pre>
<p>How can I implement such feature?</p>