model.load_state_dict(state_dict)报错出现参数与模型不匹配

项目场景：

最近在使用UNETR做医学图像的分割，当遇到如下问题：

1. 在训练的途中突然断掉需要拿之前保存的checkpoint加载后继续训练
2. 或者使用训练好的预训练模型继续训练
   就需要用到预训练的命令和代码，成为微调

---

问题描述

输入命令行如下

python main.py --feature_size 32 --batch_size 1 --logdir unetr_test --optim_lr 1e-4 --lrschedule warmup_cosine --infer_overlap 0.5 --save_checkpoint --data_dir "/data/MAE/research-contributions-main/UNETR/BTCV/dataset/dataset0/" --pretrained_dir "/data/MAE/research-contributions-main/UNETR/BTCV/pretrained_models/" --noamp --pretrained_model_name "model_final.pt" --distributed --resume_ckpt

model.load_state_dict(state_dict)报错出现参数与模型不匹配

---

原因分析：

之前保存参数时还加入了epoch loss等模型不需要的参数，

---

解决方案：

只需要在加载参数时把 model.load_state_dict(state_dict)换成 model.load_state_dict(state_dict, strict=False)即可解决

参考：

1. model.load_state_dict(state_dict, strict=False)
2. model.load_state_dict(state_dict)报错问题