Replacing nn.Conv2d with standardized convolutions but keeping weight values

agaldran: <p>Hello,</p>
<p>I am interested in taking a model with pretrained weights, and replacing all 2d convolution layers by a version of those that incorporates weight standardization <strong>without losing the pretrained values in the layer</strong>. An implementation of such a <code>StdConv2d</code> layer could be something as simple as:</p>
<pre><code class="lang-auto">class StdConv2d(nn.Conv2d):
  def forward(self, x):
    w = self.weight
    v, m = torch.var_mean(w, dim=[1, 2, 3], keepdim=True, unbiased=False)
    w = (w - m) / torch.sqrt(v + 1e-10)
    return F.conv2d(x, w, self.bias, self.stride, self.padding, 
                          self.dilation, self.groups)
</code></pre>
<p>which I got from the BiT repo <a href="https://github.com/google-research/big_transfer/blob/140de6e704fd8d61f3e5ea20ffde130b7d5fd065/bit_pytorch/models.py#L25" rel="noopener nofollow ugc">here</a>.</p>
<p>I have been following <a href="https://github.com/pytorch/vision/issues/2391#issuecomment-653900218" rel="noopener nofollow ugc">this approach</a> to replace the layers, which means that I am running:</p>
<pre><code class="lang-auto">import torch, torch.nn as nn, torchvision.models as models

def get_layer(model, name):
    layer = model
    for attr in name.split("."):
        layer = getattr(layer, attr)
    return layer

def set_layer(model, name, layer):
    try:
        attrs, name = name.rsplit(".", 1)
        model = get_layer(model, attrs)
    except ValueError:
        pass
    setattr(model, name, layer)
    
class StdConv2d(nn.Conv2d):
    def forward(self, x):
        w = self.weight
        v, m = torch.var_mean(w, dim=[1, 2, 3], keepdim=True, unbiased=False)
        w = (w - m) / torch.sqrt(v + 1e-10)
        return F.conv2d(x, w, self.bias, self.stride, 
                              self.padding, self.dilation, self.groups)
</code></pre>
<p>and then:</p>
<pre><code class="lang-auto"># let us replace conv2d layers in a resnet18; are pretrained weights kept?
model = models.resnet18(pretrained=True)
old_first_layer_weights_avg = model.conv1.weight.mean()

for name, module in model.named_modules():
    if isinstance(module, nn.Conv2d):
        # Get current bn layer
        conv = get_layer(model, name)
        # Create new in layer
        new_conv = StdConv2d(conv.in_channels, conv.out_channels, 
                             conv.kernel_size, conv.stride, conv.padding, 
                             conv.dilation, conv.groups,  conv.bias, 
                             conv.padding_mode)
        set_layer(model, name, new_conv)

new_first_layer_weight_avg = model.conv1.weight.mean()
print(old_first_layer_weights_avg == new_first_layer_weight_avg )
</code></pre>
<p>Which gives <code>false</code>, meaning that I am losing the pretrained weights here because <code>StdConv2d</code> is randomly initialized.</p>
<p>A related issue is that the above works because <code>resnet18</code> does not have conv2d layers with a bias, but when there is a bias, I get an error. This can be seen by trying to replace the conv layers with the above code in e.g. <code>alexnet</code>, which will throw the following exception:</p>
<pre><code class="lang-auto">RuntimeError: Boolean value of Tensor with more than one value is ambiguous
</code></pre>
<p>It seems to me that there must be a simple way to add standardization to conv layers while keeping the pretrained weights intact, but my anti-expertise in OOP prevents me from finding it, any help is appreciated <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>Thanks!</p>
<p>Adrian</p>

ptrblck: <p>I don’t see any code which would assign or copy the <code>.weight</code> parameter to the new conv layer in your code snippet, so I guess this part is missing.<br>
Something like:</p>
<pre><code class="lang-python">with torch.no_grad():
    new_conv.weight.copy_(conv.weight)
    if new_conv.bias:
        new_conv.bias.copy_(conv.bias)

set_layer(model, name, new_conv)
</code></pre>
<p>might work.</p>

agaldran: <p>Hi!</p>
<p>Thanks for your quick help, I appreciate it a lot. Your solution works 99%, the only thing is that <code>bias</code> seems to be overloaded in <code>nn.Conv2d</code> in a bit of a weird manner. At construction it needs to be a <code>bool</code>, but once the object is initialized, the very same <code>bias</code> becomes either <code>None</code> or a float tensor.</p>
<p>For that reason, if one tries to construct <code>new_conv</code> using <code>conv.bias</code> when <code>conv</code> originally had a bias, the code will crash because it expects a <code>bool</code>. I fixed that as follows:</p>
<pre><code class="lang-auto">for name, module in model.named_modules():
    if isinstance(module, nn.Conv2d):
        # Get current bn layer
        conv = get_layer(model, name)
        # Create new in layer, conv.bias is a tensor or None, bias needs to be bool
        b = False if conv.bias is None else True
        new_conv =StdConv2d(conv.in_channels, conv.out_channels, conv.kernel_size, 
                            conv.stride, conv.padding, conv.dilation, conv.groups, 
                             bias=b, padding_mode=conv.padding_mode)
        
        with torch.no_grad():
            new_conv.weight.copy_(conv.weight)
            if b: new_conv.bias.copy_(conv.bias)
        set_layer(model, name, new_conv)
</code></pre>
<p>Before closing, can I ask a follow-up question? I think this one is a bit more complex, if there is no easy way to achieve this, please forget it, it’s not so important (I hope).</p>
<p>In this <a href="https://github.com/joe-siyuan-qiao/WeightStandardization" rel="noopener nofollow ugc">repo</a> where they provide another implementation of weight standardization, at the bottom of the readme they recommend that one should ‘‘<em>only replace convolutional layers that are followed by normalization layers such as BN, GN, etc.</em>’’. Is there any easy modification of the above code that only applies layer replacement when this comes after a normalization layer?</p>
<p>Thanks again!!</p>

ptrblck: <p>This use case might not be trivial using the simple approach by iterating the <code>named_modules()</code> as the actual usage of the conv and normalization layers might be different than their initialization.<br>
E.g. you could first initialize all norm, then all conv layers, but could use them in a different order in the <code>forward</code>. A better approach might be to use <code>torch.fx</code> and use the actual graph to replace modules. The <a href="https://pytorch.org/docs/stable/fx.html">docs</a> would give some examples on graph manipulation.</p>