参数是一个字典,所以可以按照字典的格式
去除键值对
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](pretrained=True)
num_fc_ftr = model.fc.in_features ##重写分类输出层
model.fc = torch.nn.Linear(num_fc_ftr, num_classes)
###############加载自己训练的模型
print ("加载res34自己训练的模型")
pretrained_dict = torch.load("day_label_smooth_resnet34_model_best.pth.tar")["state_dict"]
for k,v in model.state_dict().items():
print (k)
print (v.shape)
pretrained_dict.pop('fc.weight') # 加载的参数直接删除全连接层的参数,打印名字看看
pretrained_dict.pop('fc.bias')
model.load_state_dict(pretrained_dict, strict=False)
修改 模型名字,按照名字匹配
model = DPN(num_init_features=64, k_R=96, G=32, k_sec=(3,4,20,3), inc_sec=(16,32,24,128), num_classes=1,decoder=args.decoder)
http = {'url': 'http://data.lip6.fr/cadene/pretrainedmodels/dpn92_extra-b040e4a9b.pth'}
pretrained_dict=model_zoo.load_url(http['url'])
model_dict = model.state_dict()
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}#filter out unnecessary keys
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
model = torch.nn.DataParallel(model).cuda()