RuntimeError in loading state_dict: Size mismatch and missing keys

[sharafeldinkabashi]

Hello,

I’m encountering a RuntimeError when attempting to run inference.py in the AneurysmNet repository. The error seems related to a size mismatch and missing keys in the decoder.Dense2 layer of the U_NetBN model architecture when loading the model checkpoint. Here are the details:

System and Environment
Windows 11, Python version: 3.12.7
Packages and Versions:
Brotli==1.0.9
certifi==2024.8.30
charset-normalizer==3.3.2
colorama==0.4.6
filelock==3.13.1
fsspec==2024.10.0
idna==3.7
imageio==2.36.0
Jinja2==3.1.4
lazy_loader==0.4
MarkupSafe==2.1.3
mkl_fft==1.3.10
mkl_random==1.2.7
mkl-service==2.4.0
mpmath==1.3.0
networkx==3.2.1
nibabel==5.3.2
numpy==2.1.1
packaging==24.1
pillow==10.4.0
pip==24.2
PySocks==1.7.1
PyYAML==6.0.2
requests==2.32.3
scikit-image==0.24.0
scipy==1.14.1
setuptools==72.1.0
sympy==1.13.1
tifffile==2024.9.20
torch==2.5.1
torchaudio==2.5.1
torchvision==0.20.1
tqdm==4.66.6
typing_extensions==4.11.0
urllib3==2.2.3
wheel==0.44.0
win-inet-pton==1.1.0

Error Message
Here’s the full traceback of the error:
(pytorch) C:\Users\shara\Desktop\AneurysmNet>python inference.py
C:\Users\shara\anaconda3\envs\pytorch\Lib\site-packages\torch_init_.py:1144: UserWarning: torch.set_default_tensor_type() is deprecated as of PyTorch 2.1, please use torch.set_default_dtype() and torch.set_default_device() as alternatives. (Triggered internally at C:\cb\pytorch_1000000000000\work\torch\csrc\tensor\python_tensor.cpp:434.)
_C._set_default_tensor_type(t)
UnetRI0
C:\Users\shara\Desktop\AneurysmNet\models.py:181: FutureWarning: You are using torch.load with weights_only=False (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling.
checkpoint=torch.load(path)
C:\Users\shara\anaconda3\envs\pytorch\Lib\site-packages\torch\serialization.py:1469: SourceChangeWarning: source code of class 'torch.nn.modules.activation.LeakyReLU' has changed.
warnings.warn(msg, SourceChangeWarning)
Traceback (most recent call last):
File "C:\Users\shara\Desktop\AneurysmNet\inference.py", line 74, in
Model=Segmentation(modtype,
File "C:\Users\shara\Desktop\AneurysmNet\models.py", line 152, in init
self.load(savefile,network)
File "C:\Users\shara\Desktop\AneurysmNet\models.py", line 188, in load
self.network.load_state_dict(checkpoint['model_state_dict'], strict=False)
File "C:\Users\shara\anaconda3\envs\pytorch\Lib\site-packages\torch\nn\modules\module.py", line 2584, in load_state_dict
raise RuntimeError(
RuntimeError: Error(s) in loading state_dict for U_NetBN:
Missing key(s) in state_dict: "bn.norm.weight", "bn.norm.bias", "bn.norm.running_mean", "bn.norm.running_var", "bn.conv.weight", "bn.conv.bias".
size mismatch for decoder.Dense2.conv1.conv.weight: copying a param with shape torch.Size([64, 128, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([64, 64, 3, 3, 3]).