Can i remove a layer from a pre-trained model while loading the model weights?

Lewis: <p>Hi,<br>
I am working on a problem that requires pre-training a first model at the beginning and then using this pre-trained model and fine-tuning it along with a second model. When training the first model, it requires a classification layer in order to compute a loss for it. However, I do not need my classification layer when using the pretrained model along with my second model. I only need the output  (which in my case is the hidden state of an LSTM). At the second stage, when loading the pre-trained model, if i remove the classification layer, then does PyTorch automatically ignore the weights for that classification layer and keep the rest when fine-tuning? Are the weights for each sub-module separately saved and then used as needed? This is a snippet of the code:</p>
<pre><code class="lang-auto">class LSTMD(nn.Module):
    
    def __init__(self, input_size, hidden_size, num_classes):
        super(LSTMD, self).__init__()
        self.hidde_size = hidden_size
        self.embedding = nn.Embedding(num_classes, hidden_size)
        self.lstm = nn.LSTMCell(input_size, hidden_size)
        self.classification = nn.Linear(hidden_size, num_classes)
        
    def forward(self, x):
        h,c = [torch.zeros(1, self.hidden_size), torch.zeros(1, self.hidden_size)]
        x = self.embedding(x)
        h,c = self.lstm(x, (h,c))
        out = self.classification(h)
        return out

model = LSTMD(512,512,1000)
state = {'model': model, 'model_optimizer': model_optimizer}

torch.save(state, 'saved.pth')
</code></pre>
<p>After that, I only want to load my pre-trained model but only use <code>h</code> from the LSTM output and discard the classification layer when fine-tuning the model.</p>
<pre><code class="lang-auto">class New(nn.Module):
    
    def __init__(self, checkpoint):
        super(New, self).__init__()
        checkpoint = torch.load(checkpoint)
        old_model = checkpoint['model']
        # I removed out = self.classification(h) from the old model
        modules = list(old_model.children())[:-1]
        self.new_model = nn.Sequential(*modules)

    def forward(self, x):

        out = self.new_model(x)  
        return out

</code></pre>
<p>But then I get the error: forward() takes _ positional arguments but _ were given<br>
_ depends on the arguments in my case.</p>
<p>However, when I remove the<br>
<code>self.classification = nn.Linear(hidden_size, num_classes)</code><br>
and<br>
<code>out = self.classification(h)</code><br>
from the first model, and load the checkpoint, I get no error. But i’m afraid this is not the correct way as later the model will have problem fine-tuning? The problem mainly happens when I call the forward function on the new model created by <code>nn.Sequential</code>.  Even if i keep the layers by doing:</p>
<pre><code class="lang-auto"> modules = list(old_model.children())
  self.new_model = nn.Sequential(*modules)
</code></pre>
<p>And then run the forward function, it dosen’t work. Something is happening is the sequential operation .Any clue what may be the problem<br>
Thanks!</p>

ptrblck: <p>Could you try to save the <code>state_dict</code> instead of the model and optimizer directly?<br>
Then while restoring, try to use <code>strict=False</code> in <code>.load_state_dict</code>.</p>

Lewis: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>. This worked perfectly. Apparently there were some extra keys that I didn’t take note of, so PyTorch can’t match them. However, I thought of a cleaner solution. Can’t I just freeze the classification layer when fine-tuning? It is the last layer so backprop wouldn’t be cut in the middle:</p>
<pre><code>for p in model.classification.parameters():
    p.requires_grad = False</code></pre>

ptrblck: <p>It depends on your use case and if you need the output of the last linear layer or not.<br>
Could you explain your use case a bit?</p>

Lewis: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>. Thanks for your reply!<br>
I don’t really need the output of the classification layer in the fine-tuning stage. I only need it when pre-training the first model for supervision (loss calculation) purposes. After the first model has been trained, I only want to use the output of the LSTM and not the classification layer, but I need to fine-tune it as well (fine-tune everything except the classification layer). My case is similar to when fine-tuning a pre-trained CNN. We remove the classification layer and fine-tune the other CNN layers.</p>

arkajit_bhattacharya: <p>I have a similar doubt. Is it possible to remove intermediate layers from pretrained models?</p>

ptrblck: <p>If the input and output shape is the same, you could replace these layers with an <code>nn.Identity</code> module.</p>

arkajit_bhattacharya: <p>Thanks for your response!<br>
I was trying to find if a model has a convolution layer and if present, add another custom layer infront of it.<br>
Pseudo Code<br>
for layer in self.model.children():<br>
if isinstance(layer, nn.Conv2d):<br>
<span class="hashtag">#add</span> a custom layer infront of it</p>
<p>Is this possible to achieve in pytorch?</p>

ptrblck: <p>I think there is not an easy solution to do it automatically. <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"><br>
You could abuse the <code>__getattr__</code> and <code>__setattr__</code> methods and manipulate the modules with it.<br>
E.g. here is a working (but not recommended) code snippet to add <code>nn.Identity()</code> modules in fron of all conv layers in a resnet:</p>
<pre><code class="lang-python">model = models.resnet18()

# Get all layer names you would like to change
layers_to_change = []
for name, module in model.named_modules():
    if isinstance(module, nn.Conv2d):
        print('found ', name)
        layers_to_change.append(name)

# Iterate all layers to change
for layer_name in layers_to_change:
    # Check if name is nested
    *parent, child = layer_name.split('.')
    # Nested
    if len(parent) &gt; 0:
        # Get parent modules
        m = model.__getattr__(parent[0])
        for p in parent[1:]:    
            m = m.__getattr__(p)
        # Get the conv layer
        orig_layer = m.__getattr__(child)
    else:
        m = model.__getattr__(child)
        orig_layer = copy.deepcopy(m) # deepcopy, otherwise you'll get an infinite recusrsion
    # Add your layer here
    m.__setattr__(child, nn.Sequential(
        nn.Identity(), orig_layer))

print(model)
</code></pre>
<p>While this seem to work for this model, I would recommend to create a new model class by deriving from the desired model as the base class and manipulate the layers inside the <code>__init__</code> method.</p>

arkajit_bhattacharya: <p>Thanks a lot <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
I had another doubt. Would really appreciate the help<br>
If my last layer has one output, i.e., my target is to train a model to get only one class(example: cat)…which loss function would be adequate for it? I was using MSE but it doesn’t converge. It goes past the target value if more data is used.</p>

ptrblck: <aside class="quote no-group" data-username="arkajit_bhattacharya" data-post="10" data-topic="57899">
<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/arkajit_bhattacharya/48/14907_2.png" class="avatar"> arkajit_bhattacharya:</div>
<blockquote>
<p>i.e., my target is to train a model to get only one class(example: cat)</p>
</blockquote>
</aside>
<p>I assume you are still dealing with two classes, i.e. cat vs. not-cat?<br>
If you are dealing with a single class, you cannot learn anything.</p>
<p>On the other hand, for a binary classification (cat vs. not-cat), you could use <code>nn.BCEWithLogitsLoss</code>.</p>

miranda: <p>If I have <code>nn.LSTM(input_size, hidden_size, num_layers=2,batch_first=True, bidirectional=True)</code>  in my network, is there anyway to remove the second layer of lstm after training?</p>

ptrblck: <p>Based on <a href="https://discuss.pytorch.org/t/clarification-regarding-the-return-of-nn-gru/47363/2">this</a> great visualization by <a class="mention" href="/u/vdw">@vdw</a>, you could probably try to set the <code>weight</code> and <code>bias</code> is the corresponding layer to ones and zero, and index the output state accordingly.<br>
Note that you will find the corresponding layer parameters e.g. by <code>model.weight__hh_l0</code>, where <code>l0</code> gives you the layer.</p>

miranda: <p>Thank you for your response! Do you mean lstm.bias_hh_l1, lstm.bias_ih_l1, lstm.weight_hh_l1,  lstm.weight_ih_l1, plus the reverse version of these? Because I am trying to remove the second layer of lstm in the state dictionary of the model? And can you please show it with an example for more clarification?</p>

miranda: <p>My problem is that I do not see a linear relationship between these weights and biases and the inputs and outputs of an lstm unit. So, I am confused which elements should be zero or one. Is it at all possible to do so? In case it is possible, I would be grateful if you can help me with this. Otherwise, what do you think is the best way to approach this problem?</p>

arkajit_bhattacharya: <p>Thanks for your response!<br>
The task in hand is to train the model to have a variance of one. So, I have a target class one, and I am training the model through MSELoss and Stochastic Gradient descent. The issue is I am not able to train the model. It never converges.</p>

dcbm: <p>I’m always learning from you. Thank you very much!</p>