######### multi gpu load one gpu need mudule.
print ("model load param###########################")
pretrained_dict = torch.load("model_ir_se50.pth") ###
self.model_dict = self.model.state_dict() #get the name:value
param={}
for k, v in pretrained_dict.items():
if k[7:] !="module" :
param["module."+k] = pretrained_dict[k]
pretrained_dict = { k: v for k, v in param.items() if k in self.model_dict}
self.model_dict.update(pretrained_dict)
self.model.load_state_dict(self.model_dict)
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()
简单写法
直接pop 去除全连接参数,然后加载
###############加载自己训练的模型
print ("加载自己训练的模型")
pretrained_dict = torch.load("model_best.pth.tar")["state_dict"] #保存的模型的全部结果包括全连接,优化器
for k,v in pretrained_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)
参数 打印结果,有分类输出层
_blocks.15._se_reduce.bias
torch.Size([48])
_blocks.15._se_expand.weight
torch.Size([1152, 48, 1, 1])
_blocks.15._se_expand.bias
torch.Size([1152])
_blocks.15._project_conv.weight
torch.Size([320, 1152, 1, 1])
_blocks.15._bn2.weight
torch.Size([320])
_blocks.15._bn2.bias
torch.Size([320])
_blocks.15._bn2.running_mean
torch.Size([320])
_blocks.15._bn2.running_var
torch.Size([320])
_blocks.15._bn2.num_batches_tracked
torch.Size([])
_conv_head.weight
torch.Size([1280, 320, 1, 1])
_bn1.weight
torch.Size([1280])
_bn1.bias
torch.Size([1280])
_bn1.running_mean
torch.Size([1280])
_bn1.running_var
torch.Size([1280])
_bn1.num_batches_tracked
torch.Size([])
_fc.weight
torch.Size([115, 1280])
_fc.bias
torch.Size([115])
Using image size 224