Load optimizer at specific layers?

username: Hedwig_Huang
cooked: <p>I defined two models, which is generally the same structure,<br>
while one is pre-trained, called as <strong>teacher</strong> model,<br>
the other one is initialized from the <strong>teacher</strong> model, and defined some new layer, we named <strong>student</strong> model. (but not very strictly speaking teacher-student one.)</p>
<p>In the initialilze part, we load the <strong>pretrained model’s weight</strong> as well as the <strong>optimizer state_dict</strong> .</p>
<p>As I mentioned above,  there are some difference of the structure with the teacher and the student, so when loading the pretrained optimzier state_dict…Error occurs like this:</p>
<p><code>ValueError: loaded state dict has a different number of parameter groups</code></p>
<p>But unfortunately, it seems not any solutions for loaded specific named layers with some clearly indexing like <code>net.layer1.weight</code> or some stuff to get the corresponding param_groups.</p>
<p>I found some solutions from <a href="https://blog.csdn.net/Answer3664/article/details/104874243" rel="noopener nofollow ugc">here</a>. He firstly, freeze the new layer with <code>requires_grad</code> as <code>False</code>, then loading or some stuff. I modified this solution fit my case.</p>
<p>Here is my sample code:</p>
<p><strong>Define the network</strong></p>
<pre><code class="lang-auto">import torch.nn as nn
import torch
import random

class Mynet1(nn.Module):    # teacher
    def __init__(self):
        super(Mynet1, self).__init__()
        self.layer1 = nn.Linear(1, 1)
        self.layer2 = nn.Linear(1, 1)
    
    def forward(self, x):
        out = self.layer1(x)
        out2 = self.layer2(out)

        return out2


class Mynet2(nn.Module):   # student
    def __init__(self):
        super(Mynet2, self).__init__()
        self.layer1 = nn.Linear(1, 1)
        self.layerK = nn.Linear(1, 1)
        self.layer2 = nn.Linear(1, 1)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layerK(out)   # new defined layer
        out2 = self.layer2(out)

        return out2
</code></pre>
<p><strong>Train main.py</strong></p>
<pre><code class="lang-auto"># the teacher network, pretend that is well trained and defined and loaded
net1 = Mynet1()  
opt1 = torch.optim.SGD(net1.parameters(), lr=0.01, weight_decay=3e-5, momentum=0.99, nesterov=True) 

# the student network
net2 = Mynet2()

loss = nn.L1Loss()
  
################################################
## Here is how I defined opt2 and load opt1 ##

# frozen the layerK first!
for name, param in net2.named_parameters():
    if 'K' in name:
        param.requires_grad = False
    else:
        param.requires_grad = True

opt2 = torch.optim.SGD((filter(lambda p: p.requires_grad, net2.parameters())), lr=0.01)

state = opt1.state_dict()

opt2.load_state_dict(state)
opt2.add_param_group({'params': net2.layerK.parameters(), 'lr': 0.02, 'weight_decay': 1, 'momentum': 0.99, 'nesterov': True})

# release layerK so all the layers can be trained
for name, param in net2.named_parameters():
    param.requires_grad = True

#################################################################

    # training
    net2.train()
    for i in range(500):
        x = y = random.randint(0, 10)
        x = torch.tensor([x]).type(dtype=torch.float)
        y = torch.tensor([y]).type(dtype=torch.float) 

        pred = net2(x)
        opt2.zero_grad()
        l = loss(pred, y)
        l.backward()
        opt2.step()


</code></pre>
<p>Could you please tell will the <code>net2</code> and <code>opt2</code> work as we want? like updating <code>layer1</code> and <code>layer2</code> with one settings while <code>layerK</code> with other settings I defined?<br>
Is it reasonable?</p>

username: Hedwig_Huang
cooked: <p>Hi dear sir! <a class="mention" href="/u/ptrblck">@ptrblck</a><br>
I want to modify my optimizer with some layers with <code>lr1</code> while others with <code>lr2</code>. Can I just do it with freeze some layers first and then add up some param_groups later? it is work?</p>
<p>Thank you for your helps!</p>