训练好的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(): 遍历每一层修改层名

你可能感兴趣的:(基本算法)