Change Name of State_Dict in a pretrained model

Vatsal_Malaviya: <p>I have a model following DenseNet Architecture, but the problem with is that checkpoints in pre-trained version<br>
Some of the keys for example:<br>
DenseNet121 Version:</p>
<pre><code class="lang-auto">densenet121.features.conv0.weight
densenet121.features.norm0.weight
densenet121.features.norm0.bias
densenet121.features.norm0.running_mean
densenet121.features.norm0.running_var
densenet121.features.norm0.num_batches_tracked
densenet121.features.denseblock1.denselayer1.norm1.weight
densenet121.features.denseblock1.denselayer1.norm1.bias
densenet121.features.denseblock1.denselayer1.norm1.running_mean
densenet121.features.denseblock1.denselayer1.norm1.running_var
densenet121.features.denseblock1.denselayer1.norm1.num_batches_tracked
densenet121.features.denseblock1.denselayer1.conv1.weight
densenet121.features.denseblock1.denselayer1.norm2.weight
densenet121.features.denseblock1.denselayer1.norm2.bias
densenet121.features.denseblock1.denselayer1.norm2.running_mean
densenet121.features.denseblock1.denselayer1.norm2.running_var
densenet121.features.denseblock1.denselayer1.norm2.num_batches_tracked
densenet121.features.denseblock1.denselayer1.conv2.weight
densenet121.features.denseblock1.denselayer2.norm1.weight
densenet121.features.denseblock1.denselayer2.norm1.bias
densenet121.features.denseblock1.denselayer2.norm1.running_mean
densenet121.features.denseblock1.denselayer2.norm1.running_var
densenet121.features.denseblock1.denselayer2.norm1.num_batches_tracked
densenet121.features.denseblock1.denselayer2.conv1.weight
</code></pre>
<p>Pretrained Version:</p>
<pre><code class="lang-auto">module.densenet121.features.conv0.weight
module.densenet121.features.norm0.weight
module.densenet121.features.norm0.bias
module.densenet121.features.norm0.running_mean
module.densenet121.features.norm0.running_var
module.densenet121.features.denseblock1.denselayer1.norm.1.weight
module.densenet121.features.denseblock1.denselayer1.norm.1.bias
module.densenet121.features.denseblock1.denselayer1.norm.1.running_mean
module.densenet121.features.denseblock1.denselayer1.norm.1.running_var
module.densenet121.features.denseblock1.denselayer1.conv.1.weight
module.densenet121.features.denseblock1.denselayer1.norm.2.weight
module.densenet121.features.denseblock1.denselayer1.norm.2.bias
module.densenet121.features.denseblock1.denselayer1.norm.2.running_mean
module.densenet121.features.denseblock1.denselayer1.norm.2.running_var
module.densenet121.features.denseblock1.denselayer1.conv.2.weight
module.densenet121.features.denseblock1.denselayer2.norm.1.weight
module.densenet121.features.denseblock1.denselayer2.norm.1.bias
module.densenet121.features.denseblock1.denselayer2.norm.1.running_mean
module.densenet121.features.denseblock1.denselayer2.norm.1.running_var
module.densenet121.features.denseblock1.denselayer2.conv.1.weight
module.densenet121.features.denseblock1.denselayer2.norm.2.weight
module.densenet121.features.denseblock1.denselayer2.norm.2.bias
module.densenet121.features.denseblock1.denselayer2.norm.2.running_mean
module.densenet121.features.denseblock1.denselayer2.norm.2.running_var
module.densenet121.features.denseblock1.denselayer2.conv.2.weight
module.densenet121.features.denseblock1.denselayer3.norm.1.weight
module.densenet121.features.denseblock1.denselayer3.norm.1.bias
module.densenet121.features.denseblock1.denselayer3.norm.1.running_mean
module.densenet121.features.denseblock1.denselayer3.norm.1.running_var
</code></pre>

Vatsal_Malaviya: <p>I am trying to use ChexNet model , <a href="https://drive.google.com/drive/folders/1AkBlPirkJeqVU4AlRxzo8-yTFwEQQFqA?usp=sharing" rel="nofollow noopener">Link</a> contains the kays of state_dict of the ChexNet in comparison to ResNet 121,169,201,<br>
There are few GitHub repo trying to replicate ChexNet but while I tried to use their work to get the trained model I faced the same issue in all them,<br>
Errors for Missing and Unexpected Keys while loading the state_dict</p>
<pre><code class="lang-auto">if checkpoint != None:
    modelCheckpoint = torch.load(checkpoint)
    model.load_state_dict(modelCheckpoint['state_dict'])
</code></pre>
<p>Any help is deeply appreciated !!</p>

ptrblck: <p>You could manually remove the <code>module</code> string in each key using one of <a href="https://discuss.pytorch.org/t/dataparallel-changes-parameter-names-issue-with-load-state-dict/60211/2">these approaches</a>.</p>
<p>The better approach would be to store the <code>state_dict</code> of the plain model (not the <code>nn.DataParallel</code> model) via <code>torch.save(model.module.state_dict(), PATH)</code>, which would avoid adding the <code>module</code> names.</p>

Vatsal_Malaviya: <p>Also, num_batches_tracked is and extra layer in the newer version of pytorch densenet model, therefore in the pretrained version this layer is missing. Is there a way to leave this layer and copy rest or any suggestions you might have?</p>

ptrblck: <p>You could try to use <code>strice=False</code> in <code>model.load_state_dict</code> to ignore these mismatches.</p>

s_karki: <p>I assume it is <code>strict=False</code></p>

univanxx: <p>The main task is to remove the <strong>module.</strong> part from the beginning of the parameter name. I had the similar problem and this solution worked:</p>
<pre><code class="lang-auto">state_dict = torch.load("/model_weights_path.pth")
state_dict_fixed = {}

for key in state_dict:
    # remove the module. part
    state_dict_fixed[key[7:]] = state_dict[key]

model.load_state_dict(state_dict_fixed)
</code></pre>
<p>The solutions is to create new dictionary with the same parameters but different names.</p>