pytorch加载模型报错RuntimeError: Error(s) in loading state_dict for SSD:Missing key(s) in state_dict:...

报错信息：

RuntimeError: Error(s) in loading state_dict for SSD:
Missing key(s) in state_dict: “backbone.vgg.0.weight”, “backbone.vgg.0.bias”, “backbone.vgg.2.weight”, “backbone.vgg.2.bias”, “backbone.vgg.5.weight”, “backbone.vgg.5.bias”, “backbone.vgg.7.weight”, “backbone.vgg.7.bias”, “backbone.vgg.10.weight”, “backbone.vgg.10.bias”, “backbone.vgg.12.weight”, “backbone.vgg.12.bias”,…

Unexpected key(s) in state_dict: “0.weight”, “0.bias”, “2.weight”, “2.bias”, “5.weight”, “5.bias”, “7.weight”, “7.bias”, “10.weight”, “10.bias”, “12.weight”, “12.bias”,…

原因：

加载使用模型时和训练模型时的环境不一致,可以看到丢失的Keys和没定义的keys前面相差一个backbone.vgg。

解决方法：

在ssd_model.py中加载预训练模型的语句后面，加上一个False就可以了。

出错代码：

def load_pretrained_weight(self, weight_pkl):
        self.load_state_dict(torch.load(weight_pkl))

改为：

def load_pretrained_weight(self, weight_pkl):
        self.load_state_dict(torch.load(weight_pkl),False)

---

另外，如果预训练模型是在gpu上训练的，加载到cpu上也会报错，解决办法如下：

net.load_state_dict(torch.load(model_file, map_location='cpu'))

不设置map_location参数则会报错：

RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device(‘cpu’) to map your storages to the CPU.