How to replace a layer with own custom variant?

Mxbonn: <p>I am trying to replace a certain type of layer (e.g. nn.Conv2d) with my own custom variant in a way that I do not have to copy paste the model definition code and replace the relevant lines.</p>
<p>I tried to do this using the solution mentioned in: <a href="https://discuss.pytorch.org/t/replacing-convs-modules-with-custom-convs-then-notimplementederror/17736/3" class="inline-onebox">Replacing convs modules with custom convs, then NotImplementedError</a></p>
<p>However when using this with the resnet model from torchvision and trying to replace the nn.Conv2d layers it does not replace the conv2d layers in the downsample block, and when trying to use it with the alexnet model from torchvision it does not replace any conv2d layer.</p>
<p>What is the recommended way to replace a certain type of layer with another in such a way that the original code definition of the model does not need to be copy pasted and adapted?</p>

sansmoriaxz: <p>Layers can only be replaced when the input and output dimensions are same.</p>
<p>And unless I am mistaken the layers in ResNet are present as layer blocks not individually. So either you must replace an entire block of them.</p>
<p>Anyway I think it’s better to replace the nodes near the end as they are more likely to carry less information in their weights. And the dense layers at they end. Those shouldn’t give any trouble.</p>
<p>Just be careful with dimensions and you should be fine I think.<br>
Take a preview of the model by viewing it’s repr before trying anything.</p>

Mxbonn: <p><a class="mention" href="/u/sansmoriaxz">@sansmoriaxz</a> I think you didnt fully understand my question I’m talking about replacing a layer type with another type for example, replacing torch.nn.Conv2d(*args, **kwargs) with mypackage.CustomConv2d(*args, **kwargs)</p>
<p><a class="mention" href="/u/alband">@albanD</a> you replied in a similar thread <a href="https://discuss.pytorch.org/t/how-to-replace-all-relu-activations-in-a-pretrained-network/31591/3" class="inline-onebox">How to replace all ReLU activations in a pretrained network?</a> that it can be done inplace, but it is not clear to me exactly how that should be done.</p>

albanD: <p>Hi,</p>
<p>There is no general guideline as it depends a lot on how the network you want to modify is structured as nn.Modules.<br>
If you have your convs as <code>self.conv1</code>, <code>self.conv2</code> etc, then you need to change these.<br>
If they are in a Sequential, you can find them and replace the <code>self.modules[conv_idx]</code> value for each.<br>
If it’s in the model definition in your python file, you can use another function like:</p>
<pre><code class="lang-auto">def new_conv(*args, **kwargs):
  if cfg.use_custom_conv:
    return MyConv(*args, **kwargs)
  else:
    nn.Conv2d(*args, **kwargs)

# In your net definition
self.conv = new_conv(foo, bar)
</code></pre>

Oli: <p>An easy way that might work in some cases is creating a module that contains a conv layer but also does something before or after the conv layer has done its forward pass.</p>

achyut-srivastava: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/8/18642cd3df8d09a8f7b85446bd31d19750bf7807.png" data-download-href="https://discuss.pytorch.org/uploads/default/18642cd3df8d09a8f7b85446bd31d19750bf7807" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/1/8/18642cd3df8d09a8f7b85446bd31d19750bf7807.png" alt="image" data-base62-sha1="3tM5aRa1MXn7lu0Wo4Xc9ghDR8H" width="690" height="132" data-dominant-color="ECEDEC"><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">895×172 10.3 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>you can use something like this for replacing layers.<br>
The picture above will replace all MaxPool2d layers with AvgPool2d</p>

Brando_Miranda: <p>I think this will work for you, just change it to your custom layer. Let us know if did work:</p>
<pre><code class="lang-auto">def replace_bn(module, name):
    '''
    Recursively put desired batch norm in nn.module module.

    set module = net to start code.
    '''
    # go through all attributes of module nn.module (e.g. network or layer) and put batch norms if present
    for attr_str in dir(module):
        target_attr = getattr(m, attr_str)
        if type(target_attr) == torch.nn.BatchNorm2d:
            print('replaced: ', name, attr_str)
            new_bn = torch.nn.BatchNorm2d(target_attr.num_features, target_attr.eps, target_attr.momentum, target_attr.affine,
                                          track_running_stats=False)
            setattr(module, attr_str, new_bn)

    # iterate through immediate child modules. Note, the recursion is done by our code no need to use named_modules()
    for name, immediate_child_module in module.named_children():
        replace_bn(immediate_child_module, name)

replace_bn(model, 'model')
</code></pre>
<p>original post: <a href="https://discuss.pytorch.org/t/how-to-modify-a-pretrained-model/60509/10" class="inline-onebox">How to modify a pretrained model</a></p>

nima_rafiee: <p>I tried this for resnet to change BN into GN. but it keeps 3 BN unchanged.</p>
<pre><code class="lang-auto">  (layer1): Sequential(
    (0): Bottleneck(
      (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn1): GroupNorm(1, 64, eps=1e-05, affine=False)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): GroupNorm(1, 64, eps=1e-05, affine=False)
      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn3): GroupNorm(1, 256, eps=1e-05, affine=False)
      (relu): ELU(alpha=1.0, inplace=True)
      (downsample): Sequential(
        (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
</code></pre>

Antiger: <p>Nice tutorial! It is a very informative tutorial. It was extremely helpful<a href="https://www.tellpizzahut.online/" rel="noopener nofollow ugc">.</a></p>

Samue1: <p>I get the following error when I try this approach:</p>
<p><code>TypeError: 'VGG' object is not iterable</code></p>
<p>Any idea why this is the case?</p>

Samue1: <aside class="quote no-group" data-username="Brando_Miranda" data-post="7" data-topic="43586">
<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/brando_miranda/48/14355_2.png" class="avatar"> Brando_Miranda:</div>
<blockquote>
<p><code>target_attr = getattr(m, attr_str)</code></p>
</blockquote>
</aside>
<p>I get an error since <code>m</code> is not defined. Any idea what I do wrong? Where does <code>m</code> come from?</p>

Qwwq: <p>Sometimes when the layers are in a list, then it creates a problem with getattr as it does not understand something like ‘layer1.0.bn1’ which you would otherwise address as net.layer1[0].bn1. To resolve this I have created the following script.</p>
<pre><code class="lang-auto">def get_all_parent_layers(net, type):
    layers = []

    for name, l in net.named_modules():
        if isinstance(l, type):
            tokens = name.strip().split('.')

            layer = net
            for t in tokens[:-1]:
                if not t.isnumeric():
                    layer = getattr(layer, t)
                else:
                    layer = layer[int(t)]

            layers.append([layer, tokens[-1]])

    return layers


def set_all_bn_to_identity(net):
    for parent_layer, last_token in get_all_parent_layers(net, nn.BatchNorm2d):
        setattr(parent_layer, last_token, nn.Identity())
</code></pre>
<p>First we return all the parent layers and the last tokens of each parent layer that satisfies a particular layer type such as nn.BatchNorm2d. Then we can simply iterate over the list of all (parent_layer, last_token) and use setattr to perform our desired change.</p>

Sieun_Park: <pre><code class="lang-auto">import timm 
import torch
model = timm.create_model('resnet34')

getattr(model.layer1, "0")
&gt;&gt;&gt; BasicBlock(
  (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (drop_block): Identity()
  (act1): ReLU(inplace=True)
  (aa): Identity()
  (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (act2): ReLU(inplace=True)
)

model.layer1[0]
&gt;&gt;&gt; BasicBlock(
  (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (drop_block): Identity()
  (act1): ReLU(inplace=True)
  (aa): Identity()
  (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (act2): ReLU(inplace=True)
)

model.layer1.0
&gt;&gt;&gt;  File "&lt;ipython-input-33-2ff1b8dc6a47&gt;", line 1
    model.layer1.0
                 ^
SyntaxError: invalid syntax
</code></pre>
<p><a class="mention" href="/u/qwwq">@Qwwq</a> Huge thank you for your solution, <code>model.named_modules</code> was what I was looking for and I really like your code!</p>
<p>However, it seems to me that <code>getattr</code> understands addresses such as <code>net.layer1[0].bn1</code> as seen in the example below? I also observed the same using <code>torchvision</code>’s ResNet. Nothing serious though:)</p>

Qwwq: <p>That is correct. <code>getattr</code> (and <code>setattr</code>) does understand <code>net.layer1[0]</code>, and you can <em>manually</em> resolve pytorch module names such as <code>net.layer1.0.bn1</code> to <code>net.layer1[0].bn1</code>. The problem occurs if you have too many such resolutions to make, which is say when you want to replace <em>all</em> BN layers in a very large network to something else. If you were to address each module name individually, it becomes tedious enough to warrant an automated solution.</p>

phucdoitoan: <aside class="quote quote-modified" data-post="16" data-topic="60068">
  <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/phucdoitoan/48/68986_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-replace-a-layer-or-module-in-a-pretrained-network/60068/16">How to Replace a layer or Module in a pretrained network?</a> 
  </div>
  <blockquote>
    Just in case people find this useful, you can replace specific layers in a pretrained network with your customed layer iteratively as follow (or modify it according to your need). 
def replace_module(module, target_name, new_module):
    for child_name, child_module in module.named_children():
        if target_name in child_name:
            setattr(module, child_name, new_module)
        else:
            replace_module(child_module, target_name, new_module)

# example:
t5_model = T5ForConditi…
  </blockquote>
</aside>