State_dict when it comes to test a model that is trained based on loading checkpoints

Tm4: <p>I trained a model using PyTorch and I stored the weights in which it had the minimum validation loss during training. Also, I stored the optimizer weights on that time. So, I trained my model for the second time by loading the weights that I stored from the first time. I wanted to test my model and generate the outputs of my model. But, unfortunately, I faced with such an error:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/b/7b4f7cf0780d28d1bded10b69190d1640e2eb242.png" data-download-href="https://discuss.pytorch.org/uploads/default/7b4f7cf0780d28d1bded10b69190d1640e2eb242" title="Capture"><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/b/7b4f7cf0780d28d1bded10b69190d1640e2eb242.png" alt="Capture" data-base62-sha1="hAR1VZk5lRAfrjBwkhsuKSlI7my" width="690" height="122" data-dominant-color="DAC9C9"><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">Capture</span><span class="informations">1417×252 21.9 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>When I wanted to load the weights for the second_time training I used the following scripts:</p>
<pre><code class="lang-auto">    model = nn.DataParallel(model, device_ids = [0,1],output_device= [0,1])
    
    model = model.to(dev)

    model.load_state_dict(torch.load(path_to_the_saved_model's_weights, map_location=dev))
    
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=2e-7)

    
    # loading optimizer
    optimizer.load_state_dict(torch.load(path to the saved optimizer's weights))
</code></pre>
<p>I only changed learning rate value during the second_time training and when I want to test the model and generating the output of the model by loading the new weights I see such an error. How can I solve the problem?</p>

Karthik_Ganesan: <p>This error is because the items in your model’s state_dict don’t match the state_dict you are trying to load. Just try to print out all the items in your model’s state_dict and the state_dict you get using <code>torch.load</code>. If they are not the same, you will get this error.</p>
<p>When you trained the model, did you use <code>nn.DataParallel</code>?</p>

saluei: <p>This is due to mismatch in model layer name during first training and second time.<br>
just check the name of key value in state_dic saved in the first time training and compare it with the name of state_dic in second time. if there is any mismatch in key name this code may help. just rename keys in state_dics to current name in model.</p>
<pre><code class="lang-auto">def LoadStateDictCustom(self,StateDicPath):
        StatDic=torch.load(StateDicPath,map_location=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
        StatDic=StatDic['model_state_dict']
        StatDic2=OrderedDict()
        for key,value in StatDic.items():
            if(key=='OLD key name that must be changed  '):
                StatDic2['NEW key name for second time']=value
            else:
                StatDic2[key]=value
        self.load_state_dict(StatDic2)
</code></pre>

Tm4: <p>Thans for your answer, could you please provide an example on how to use your function and also how the output of your solution should be used?</p>

saluei: <p>you could copy this function to your model class, pass StateDic Path as parameter to it,<br>
this function just load the given StateDic and create new StateDic that is same as the first one, except rename some key value to the new key value that is proper to your model</p>
<pre><code class="lang-auto">class YourModel(nn.Module):
    def __init__(self, n_classes, in_chann):
        super(YourModel, self).__init__()


    def LoadStateDictCustom(self,StateDicPath):
        StatDic=torch.load(StateDicPath,map_location=torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
        StatDic2=OrderedDict()
        for key,value in StatDic.items():
            if(key=='OLD key name that must be changed  '):
                StatDic2['NEW key name for second time']=value
            else:
                StatDic2[key]=value
        self.load_state_dict(StatDic2)

    def forward(self, x,trg=None):
         pass
</code></pre>

Tm4: <p>Thanks a lot for your update but my main question is that I have several sequential convolutional layers in the  forward function of my model’s class. In which section of the forward function I should call <code>LoadStateDictCustom</code> ? Should I call it in my script that is for training or I should call it inside the forward?</p>

saluei: <p>you should call it instead of this line</p>
<aside class="quote no-group" data-username="Tm4" data-post="1" data-topic="167236">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/b38774/48.png" class="avatar"> Tm4:</div>
<blockquote>
<p><code>model.load_state_dict(torch.load(path_to_the_saved_model's_weights, map_location=dev))</code></p>
</blockquote>
</aside>

Tm4: <p><a class="mention" href="/u/saluei">@saluei</a>  , Thank you again for your solution. I faced a new challenge. I want to load the weights and fine-tune my model while I changed some of the layers of my model. My plan is to freeze all the previous model parameters and only train specific layer. But according to this scenario, when I want to load weights I face the error in loading state_dict. How can I apply <code>strict=False</code> in your proposed solution.</p>