在测试该作者的去雾网络的预训练模型,代码链接:https://github.com/yangyiliu21/ntire2023_ITBdehaze
出现如下问题:
RuntimeError: Error(s) in loading state_dict for DataParallel:
Missing key(s) in state_dict: "module.feature_extract.encoder. patch_ embed.proj.weight"
加载网络预训练模型失败,根据项目issue板块的作者回答:
打印了模型的keys,发现模型keys=(epoch, model_state_dict)。
原有加载代码为:
MyEnsembleNet.load_state_dict(torch.load(args.ckpt_path))
修改后:
MyEnsembleNet.module.load_state_dict(torch.load(args.ckpt_path)['model_state_dict'])
模型加载成功。
总结
发现模型加载失败出现missing key问题,最好先打印一下模型的key。