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)报错问题

你可能感兴趣的:(深度学习,人工智能,机器学习)