问题描述:
原作者的代码里面没有断点续训,我增加了这个功能,同时也引入了更多的参数,保存模型时增加了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