How do I update old key name in state_dict to new version

kaltu: <p>I recently dug out an archived densenet121 weight from 2017.<br>
And now I tried to resurrect the weight using PyTorch 1.4.0<br>
It seems the naming convention had been changed due to some reasons.<br>
When I load the same architecture under newer PyTorch version with legacy state_dict, encountered the following error</p>
<pre><code class="lang-auto">RuntimeError: Error(s) in loading state_dict for DataParallel:
	Missing key(s) in state_dict: "module.densenet121.features.denseblock1.denselayer1.norm1.weight", "module.densenet121.features.denseblock1.denselayer1.norm1.bias",  ......
	Unexpected key(s) in state_dict: "module.densenet121.features.denseblock1.denselayer1.norm.1.weight", "module.densenet121.features.denseblock1.denselayer1.norm.1.bias", ......
</code></pre>
<p>Clearly somehow back in time the keys were in <code>layer.#.param</code> format and probably due to consistency that format was changed to <code>layer#.param</code></p>
<p>Is there any existing tool to automatically and generally update the old key names to the newer version?</p>

ptrblck: <aside class="quote no-group" data-username="kaltu" data-post="1" data-topic="76026">
<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/kaltu/48/22596_2.png" class="avatar"> kaltu:</div>
<blockquote>
<p>Is there any existing tool to automatically and generally update the old key names to the newer version?</p>
</blockquote>
</aside>
<p>I’m not aware of such a tool, but I think it should be easy to remove the unwanted indexing from all <code>keys</code> in the <code>dict</code>. Are you trying to only remove these index numbers or do you have any other mismatches?</p>

kaltu: <p>Indeed it is easy enough to manually change keys<br>
In my case it was done by</p>
<pre><code class="lang-auto">for key in list(state_dict.keys()):
    state_dict[key.replace('.1.', '1.'). replace('.2.', '2.')] = state_dict.pop(key)
</code></pre>
<p>It’s not generalizable, but if there is no one encountered this problem then it probably isn’t a big deal.</p>

stormchase143: <p>Hello, I have the same problem.  I had already trained a model, but now I have changed the names of the layers. After this above description, how do I then equate it to the model’s state_dict?<br>
Thank you in advance for your help!</p>

Ajinkya.Bankar: <p>Hello <a class="mention" href="/u/ptrblck">@ptrblck</a> and <a class="mention" href="/u/kaltu">@kaltu</a> , I am facing similar problem. I want to change the names of two keys in the model. I tried the suggested method, but I am getting errors. My code snippet and error message are below:</p>
<pre><code class="lang-auto">class LeNet(nn.Module):
    def __init__(self, num_classes=43, input_channels=3):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(input_channels, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1   = nn.Linear(16*5*5, 120)
        self.fc2   = nn.Linear(120, 84)
        self.fc3   = nn.Linear(84, num_classes)
        if 1 == num_classes: 
            # compatible with nn.BCELoss
            self.softmax = nn.Sigmoid()
        else:    
            # compatible with nn.CrossEntropyLoss
            self.softmax = nn.LogSoftmax(dim=1)        


    def forward(self, x):
        out = F.relu(self.conv1(x))
        out = F.max_pool2d(out, 2)
        out = F.relu(self.conv2(out))
        out = F.max_pool2d(out, 2)
        out = out.view(out.size(0), -1)
        out = F.relu(self.fc1(out))
        out = F.relu(self.fc2(out))
        out = self.fc3(out)

        out = self.softmax(out)
        return out

teacher_model = LeNet() # get the model
checkpoint = torch.load('model_best.pth.tar', map_location=device)
state_dict = checkpoint['state_dict']
for key in list(state_dict.keys()):
    state_dict[key.replace("conv1.weight","features.0.weight"). replace("conv1.bias", "features.0.bias")] = state_dict.pop(key)
teacher_model.load_state_dict(checkpoint['state_dict'])
</code></pre>
<p><strong>It is giving me the following error:</strong></p>
<pre><code class="lang-auto">Error(s) in loading state_dict for LeNet:
Missing key(s) in state_dict: "conv1.weight", "conv1.bias".
Unexpected key(s) in state_dict: "features.0.weight", "features.0.bias".
</code></pre>

Ajinkya.Bankar: <p>Hi, were you able to equate modified state_dict to the model and load it successfully? Thanks.</p>

ptrblck: <p>Based on the error message it seems that you are replacing the needed <code>conv1</code> keys with unexpected <code>features</code> keys. What kind of error were you seeing before trying to manipulate it?</p>

Ajinkya.Bankar: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thank you for the response. Initially, I had no errors and I was able to load the model which has old keys. But, to use this model in the energy calculation framework, it requires the key names as <code>"features.0.weight", "features.0.bias"</code> instead of <code>"conv1.weight”, “conv1.bias”</code> respectively. Therefore, I am trying to modify the key names and then load the model. Please help to solve the problem.</p>

ptrblck: <p>I don’t quite understand the issue, since your manipulation is creating the mismatches.<br>
Did you maybe forget to modify the model before using <code>load_state_dict</code>?</p>

Ajinkya.Bankar: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I updated the snippet with my model’s class. Do I need to modify the layer names in model class also in order to use new key names?</p>

ptrblck: <p>If your <code>state_dict</code> contains parameters stored as <code>features</code>, which should be loaded in the current <code>conv</code> layers, then yes: change the layer names in the model or alternatively the keys of the <code>state_dict</code>.</p>
<p>Based on your currently posted code your model uses the <code>conv</code> layers, the <code>state_dict</code> seems to use the same keys, you are changing them to <code>features</code>, and run into the error.</p>

Ajinkya.Bankar: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Your answer helped me to solve the problem. Thank you very much!</p>

ali1z: <p>Please include a complete code on how to load a “checkpoint.pth.tar” file and then use your code to change the key like you mentioned so we can  use model.load_state_dict(torch.load(PATH)) for more prediction</p>
<p>thanks</p>

Minxiangliu: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a> ,<br>
Can you save me?<br>
After I changed the architecture of the last layer, the key disappeared, I don’t know what went wrong…</p>
<pre><code class="lang-auto">self.model.conv_seg = nn.Sequential(
            nn.AdaptiveMaxPool3d(output_size=(1, 1, 1)),
            nn.Flatten(start_dim=1),
            nn.Dropout(0.1))
print(self.model.conv_seg)
print(self.model.conv_seg.state_dict().keys())
</code></pre>
<p>Results:</p>
<pre><code class="lang-auto">Sequential(
  (0): AdaptiveMaxPool3d(output_size=(1, 1, 1))
  (1): Flatten(start_dim=1, end_dim=-1)
  (2): Dropout(p=0.1, inplace=False)
)
odict_keys([])
</code></pre>

ptrblck: <p>None of the newly created layers contains any parameters or buffers, so the <code>state_dict</code> would be empty.</p>

Minxiangliu: <p>Hi, <a class="mention" href="/u/ptrblck">@ptrblck</a><br>
How can I make the new layer generate buffers?</p>

ptrblck: <p>Neither pooling layers, nor <code>Flatten</code> or <code>Dropout</code> need any parameters or buffers, so I’m unsure how or why you would like to add buffers to them.<br>
There is “nothing” to store in the <code>state_dict</code> for these layers <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>