训练好的Pytorch模型在导入测试

Missing key(s) in state_dict: “_conv_stem.weight”, “_bn0.weight”
Unexpected key(s) in state_dict: “module._conv_stem.weight”, “module._bn0.weight”,

生成的文件层名中多了module.前缀,原因是训练的时候你调用了model = torch.nn.DataParallel(model).cuda()

采取多GPU训练模型

解决方案1: 测试的时候加入model = torch.nn.DataParallel(model).cuda()
将定义的网络结构层名前面加上module.

解决方案2:for key in state_dict.keys(): 遍历每一层修改层名

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章