最近想着修改网络的预训练模型vgg.pth,但是发现当我加载预训练模型权重到新建的模型并保存之后。在我使用新赋值的网络模型时出现了key不匹配的问题。
#加载后保存(未修改网络)
base_weights = torch.load(args.save_folder + args.basenet)
ssd_net.vgg.load_state_dict(base_weights)
torch.save(ssd_net.state_dict(), args.save_folder + 'ssd_base' + '.pth')
# 将新保存的网络代替之前的预训练模型
ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
net = ssd_net
...
if args.resume:
...
else:
base_weights = torch.load(args.save_folder + args.basenet)
#args.basenet为ssd_base.pth
print('Loading base network...')
ssd_net.vgg.load_state_dict(base_weights)
此时会如下出错误:
Loading base network…
Traceback (most recent call last):
File “train.py”, line 264, in
train()
File “