How to delete some layers in pretrained model?

username: SungmanHong
cooked: <p>I have mobilenetV3Small Model</p>
<pre><code class="lang-auto">MobileNetV3(
  (features): Sequential(
    (0): ConvBNActivation(
      (0): Conv2d(3, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
      (1): BatchNorm2d(16, eps=0.001, momentum=0.01, affine=True, track_running_stats=True)
      (2): Hardswish()
    )
    (1): InvertedResidual(
      (block): Sequential(
        (0): ConvBNActivation(
          (0): Conv2d(16, 16, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), groups=16, bias=False)
          (1): BatchNorm2d(16, eps=0.001, momentum=0.01, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
</code></pre>
<p>But I want to remove BatchNorm2D layer only… in the pretrained model.<br>
And Save as New model as mine.<br>
How can I…?</p>

username: ptrblck
cooked: <p>There are several options to remove/replace this layer depending on your use case.</p>
<p>You could copy the source code of the model, remove the layer, change the <code>forward</code> method, and store this new model definition as a custom model. This would allow you to load this new model in your scripts and just train it if needed. The <code>state_dict</code> would not contain this layer (and in particular its parameters) anymore and you would not be able to directly load a pretrained <code>state_dict</code> anymore (using <code>strict=False</code> as a workaround should work).</p>
<p>Alternatively, you could also load the original model and replace the unwanted layer with an <code>nn.Identity</code>. While this approach might be simpler, you would need to replace this layer after each model creation.</p>

username: SungmanHong
cooked: <p>Thank you! I found source code of MobileNetV3Small and<br>
just removed some layer, it also works well.<br>
However, I can’t use pretrained network and it has overfitting problem in the validation data set…</p>
<p>==&gt; As you mentioned, when loading the pretrained model at <code>load_state_dict</code>,<br>
using <code>strict=False</code> removed unwanted layer as Magic!!! Thank you!!! Pytorch is very Nice tool!!<br>
Thank you!!! Thank you for saving my life!!! Thank you!!!</p>

username: sio277
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Replacing the unwanted layer with <code>nn.Sequential()</code> will do the same thing as the <code>nn.Identity()</code> without any side effects?</p>

username: ptrblck
cooked: <p>Yes, it would work at the moment, but I would stick to the <code>nn.Identity</code> workflow as it seems to be cleaner than depending on an empty <code>nn.Sequential</code> container to return the input activation.</p>

username: sio277
cooked: <p>Oh I see. Thank you🙂</p>

username: sio277
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Could I ask more questions about replacing the layer?</p>
<p>Let me assume a simple model:</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn


class Foo(nn.Module):
    def __init__(self):
        super(Foo, self).__init__()
        self.bar = nn.Conv2d(1, 2, 3)
        
foo = Foo()
foo.bar = nn.Linear(4, 5)  # replacement
</code></pre>
<ol>
<li>
<p>When replacing the <code>foo.bar</code> layer using assignment, do we need the guard of <code>torch.no_grad()</code> ? I’m concerned it as the new parameters are coming and the old ones are gone. Or we don’t need it as it is just a replacement of the namespace (<code>foo.bar.__dict__</code>)? When replacing the parameter itself, we needed <code>torch.no_grad()</code> as you mentioned in here: <a href="https://discuss.pytorch.org/t/how-to-assign-an-arbitrary-tensor-to-models-parameter/44082/9" class="inline-onebox">How to assign an arbitrary tensor to model's parameter? - #9 by ptrblck</a>.</p>
</li>
<li>
<p>Are the parameters inside the old layer (here the convolutional layer) freed from memory (because they can no longer be pointed to)?</p>
</li>
</ol>

username: ptrblck
cooked: <ol>
<li>
<p>I would wrap it in a <code>no_grad()</code> block, to make sure that this assignment is not tracked by Autograd.</p>
</li>
<li>
<p>If these parameters are not referenced anymore, they’ll be freed. E.g. if you’ve passed them to an optimizer before replacing them, they might be kept alive (and the new parameters would have to be added to the optimizer).</p>
</li>
</ol>

username: sio277
cooked: <p>Thank you very much for your kind explanation <img src="https://discuss.pytorch.org/images/emoji/apple/grinning_face_with_smiling_eyes.png?v=9" title=":grinning_face_with_smiling_eyes:" class="emoji" alt=":grinning_face_with_smiling_eyes:"></p>