RuntimeError: Error(s) in loading state_dict for DynamicUnet: Missing key(s) in state_dict: "layers.0.4.0.conv3.weight" | size mismatch for layers

Goal:

Pickled model and exported weights come from a separate training environment. Here, I aim to load the model and weights to run inference with new datasets.

Versions:

  • torch==1.7.1
  • fastai==2.7.7
  • fastcore==1.5.6
  • torch==1.7.1
  • torchvision==0.8.2

Code:

from fastai.vision.all import *

learn = load_learner('export.pkl', cpu=True)
learn.load('model_3C_34_CELW_V_1.1')

Traceback:

(venv) me@ubuntu-pcs:~/PycharmProjects/project$ python3 model/Run_model.py 
Traceback (most recent call last):
  File "/home/me/PycharmProjects/project/model/Run_model.py", line 4, in <module>
    learn.load('model_3C_34_CELW_V_1.1')
  File "/home/me/miniconda3/envs/venv/lib/python3.9/site-packages/fastai/learner.py", line 387, in load
    load_model(file, self.model, self.opt, device=device, **kwargs)
  File "/home/me/miniconda3/envs/venv/lib/python3.9/site-packages/fastai/learner.py", line 54, in load_model
    get_model(model).load_state_dict(model_state, strict=strict)
  File "/home/me/miniconda3/envs/venv/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1051, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for DynamicUnet:
        Missing key(s) in state_dict: "layers.0.4.0.conv3.weight", "layers.0.4.0.bn3.weight", "layers.0.4.0.bn3.bias", "layers.0.4.0.bn3.running_mean",

        size mismatch for layers.12.0.weight: copying a param with shape torch.Size([3, 99, 1, 1]) from checkpoint, the shape in current model is torch.Size([3, 291, 1, 1]).

You get this error when you change the number of classes, maybe while retraining your U net model with different number of classes.
I suggest you change the new model instance by using num_of_classes=99
The model whose pretrained weights you are trying to load has 291 classes whilst yours has 99 only.

2 Likes

Yes; I needed an updated .pkl files that worked with the weights .pth file. Thanks

1 Like