[SOLVED] Unexpected key(s) in state_dict: batches_tracked"

username: Narayana_Swamy
cooked: <p>I am getting this error when I am loading the weights for a model I trained on a GPU.  I am trying to load the model on a CPU.  Not sure what this error means.  I did check the checkpoint[‘state_dict’] keys and it does contain this ‘num_batches_tracked’ key for all the batch normalization layers.  I am using the same model on the CPU as I trained on the GPU.  I am training a UNet model.</p>
<p>checkpoint = torch.load(‘model_best.pth-zeromask-notconsider.tar’,map_location=lambda storage, loc: storage)<br>
train_conv.model.load_state_dict(checkpoint[‘state_dict’])</p>
<p>RuntimeError                              Traceback (most recent call last)<br>
 in ()<br>
----&gt; 1 train_conv.model.load_state_dict(checkpoint[‘state_dict’])</p>
<p>~\AppData\Local\conda\conda\envs\myfastai1\lib\site-packages\torch\nn\modules\module.py in load_state_dict(self, state_dict, strict)<br>
719         if len(error_msgs) &gt; 0:<br>
720             raise RuntimeError(‘Error(s) in loading state_dict for {}:\n\t{}’.format(<br>
–&gt; 721                                self.<strong>class</strong>.<strong>name</strong>, “\n\t”.join(error_msgs)))<br>
722<br>
723     def parameters(self):</p>
<p>RuntimeError: Error(s) in loading state_dict for Unet34:<br>
Unexpected key(s) in state_dict: “rn.1.num_batches_tracked”, “rn.4.0.bn1.num_batches_tracked”, “rn.4.0.bn2.num_batches_tracked”, “rn.4.1.bn1.num_batches_tracked”, “rn.4.1.bn2.num_batches_tracked”, “rn.4.2.bn1.num_batches_tracked”, “rn.4.2.bn2.num_batches_tracked”, “rn.5.0.bn1.num_batches_tracked”, “rn.5.0.bn2.num_batches_tracked”, “rn.5.0.downsample.1.num_batches_tracked”, “rn.5.1.bn1.num_batches_tracked”, “rn.5.1.bn2.num_batches_tracked”, “rn.5.2.bn1.num_batches_tracked”, “rn.5.2.bn2.num_batches_tracked”, “rn.5.3.bn1.num_batches_tracked”, “rn.5.3.bn2.num_batches_tracked”, “rn.6.0.bn1.num_batches_tracked”, “rn.6.0.bn2.num_batches_tracked”, “rn.6.0.downsample.1.num_batches_tracked”, “rn.6.1.bn1.num_batches_tracked”, “rn.6.1.bn2.num_batches_tracked”, “rn.6.2.bn1.num_batches_tracked”, “rn.6.2.bn2.num_batches_tracked”, “rn.6.3.bn1.num_batches_tracked”, “rn.6.3.bn2.num_batches_tracked”, “rn.6.4.bn1.num_batches_tracked”, “rn.6.4.bn2.num_batches_tracked”, “rn.6.5.bn1.num_batches_tracked”, “rn.6.5.bn2.num_batches_tracked”, “rn.7.0.bn1.num_batches_tracked”, “rn.7.0.bn2.num_batches_tracked”, “rn.7.0.downsample.1.num_batches_tracked”, “rn.7.1.bn1.num_batches_tracked”, “rn.7.1.bn2.num_batches_tracked”, “rn.7.2.bn1.num_batches_tracked”, “rn.7.2.bn2.num_batches_tracked”, “up1.bn.num_batches_tracked”, “up2.bn.num_batches_tracked”, “up3.bn.num_batches_tracked”, “up4.bn.num_batches_tracked”.</p>

username: Narayana_Swamy
cooked: <p>Looks like it was a version problem - I trained the model in 0.4.1 and tried to load the dict in 0.4.0.  Once I upgraded the pytorch version, everything worked fine.</p>

username: ayumiymk
cooked: <p>Hi, I meet the same question. I have some models in 0.4.1, at the same time some models in 0.4.0. And my pytorch is 0.4.0. As you say, after upgrading the pytorch version, pytorch can work fine loading 0.4.1’s model. But how about loading 0.4.0’s model?</p>
<p>Thanks !</p>

username: Xuxudong166
cooked: <p>I have the same problem with you, my model and pytorch is 0.4.1, but the problem is not solved, the same error still occurs.</p>

username: Zeeshan_Anjum
cooked: <p>I trained my model in pytorch 0.4.1, in classroom I have 0.4.0 and I can’t upgrade classroom version. it is possible that I could load 0.4.1 model in 0.4.0 or I have to downgrade my pytorch version?<br>
Thanks.</p>

username: TillChima
cooked: <p>Are there any updates or solutions yet?<br>
I am having a similar trouble loading a model.</p>
<p>Edit: I don’t think that this should be marked as solved, since the solution refers to a pytorch version imcompatibility. This sounds more like a workaround than a solution. To me it should be compatible with multiple versions of pytorch since not everybody can up- or downgrade their version.</p>

username: dattran2346
cooked: <p>You can check pytorch <a href="https://github.com/pytorch/vision/blob/1f8f7ea71c8807408a7860b8ae9ca63177975b64/torchvision/models/densenet.py#L120-L144" rel="nofollow noopener">densenet</a> for the solution. The internal <code>state_dict</code> have changed from something like <code>features.denseblock1.denselayer1.norm.1.weight</code> to <code>features.denseblock1.denselayer1.norm1.weight</code></p>
<p>Also, you can write a simple function that load the model, modify <code>state_dict</code> and save the model.</p>
<pre><code class="lang-auto">def modify_state_dict(state_dict):
    pattern = re.compile(
                r'^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$')
    state_dict = checkpoint['state_dict']
    for key in list(state_dict.keys()):
        res = pattern.match(key)
        if res:
            new_key = res.group(1) + res.group(2)
            state_dict[new_key] = state_dict[key]
            del state_dict[key]

checkpoint = torch.load(model_path)
modify_state_dict(checkpoint['state_dict'])
torch.save(checkpoint, model_path)
</code></pre>
<p>For the <code>num_batches_tracked</code>, pytorch has added in later version. I have checked the value of these key in densenet layer and they are all <code>tensor(0, device='cuda:0')</code>. I think you can add that missing key to the <code>state_dict</code> when modifying it. Btw, I load my trained densenet model on v0.4.1 without that <code>num_batches_tracked</code> key.</p>

username: Y_Cheng
cooked: <p>Agree with you!!Just downgrade the version to avoid this error is not elegant.</p>