模型加载出现问题，RuntimeError: Error(s) in loading state_dict for UNet：Missing key(s) in state_dict...

问题：

模型在训练过程中可以正常训练，但是测试的时候出现了错误，如下所示：

原因分析：
可以看到下面的Unexpected key(s) in state_dict: "module.encoder.in_conv.conv_conv.0.weight"比上面的Missing key(s) in state_dict: "encoder.in_conv.conv_conv.0.weight"多了关键字module，这是因为用DataParallel训练的模型数据并行方式的，key中会包含”module“关键字。

解决方案：
(1)直接把要加载参数的模型也放在DataParallel上面，如下：

model = ResNet50().cuda()
model = nn.DataParallel(model)
checkpoint = torch.load(args.net_cache)
model.load_state_dict(checkpoint['state_dict']`

（2）去掉DataParallel 预训练model中的module(可行权重值一致)

model = ResNet50().cuda()
checkpoint = torch.load(args.net_cache)
model.load_state_dict({k.replace('module.',''):v for k,v in checkpoint['state_dict'].items()})

参考链接