PyTorch加载模型时报错RuntimeError: Error(s) in loading state_dict for *****: Missing key(s) in state_dict:

问题描述：

    原作者的代码里面没有断点续训，我增加了这个功能，同时也引入了更多的参数，保存模型时增加了epoch、net.state_dict()、optimizer.state_dict()、scheduler.state_dict()等信息。

原来的保存模型的代码如下：

torch.save(net.state_dict(), model_dir)

增加了信息后，保存模型的代码如下：

 torch.save({'epoch': i,
             'model_state_dict': net.state_dict(),
             'optimizer_state_dict': optimizer.state_dict(),
             'scheduler': scheduler.state_dict(),},
              model_dir)

原来加载模型的代码如下：

net.load_state_dict(torch.load(model_dir))

增加了信息后，加载模型的代码如下：

ckpt = torch.load(model_dir, map_location='cpu')
net.load_state_dict(ckpt['model_state_dict'])

测试推理的时候，加载模型报错：

net.load_state_dict(ckpt['model_state_dict'])
  File "/root/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1223, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for *****:
	Missing key(s) in state_dict: 

解决办法：

方法1、可以加载成功，但会导致一些参数加载不进来，某些情况下会造成推理结果错误。

ckpt = torch.load(model_dir)
model.load_state_dict(ckpt['model_state_dict'],strict=False)

方法2、将字典键值中的module.替换掉，或把原模型pth文件的key打印与现在模型的key进行比较，手动的为模型加载参数。

ckpt = torch.load(args.weights, map_location='cpu')
net.load_state_dict({k.replace('module.', ''): v for k, v in ckpt['model_state_dict'].items()})

问题根源：

训练时的代码里加入了如下代码：

net = nn.DataParallel(net)

找到训练代码中的net = nn.DataParallel(net)，注释掉再重新训练。

或者采用上述方法2去加载模型。

多GPU的并行计算，训练同一个模型可以用上Pytorch的nn.DataParallel