pytorch pretrained model
two methods
method 1 比较大小
self._model = models.__dict__[args.arch](pretrained = False,
num_classes = args.classes,
aux_logits = False)
if self.args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
pretrained_state = model_zoo.load_url(model_names[args.arch])
model_state = self._model.state_dict()
pretrained_state = { k:v for k,v in pretrained_state.iteritems() if k in model_state and v.size() == model_state[k].size() }
model_state.update(pretrained_state)
self._model.load_state_dict(model_state)
method 2 比较key
given:
pretrained_dict: ['A', 'B', 'C', 'D']
model_dict: ['A', 'B', 'C', 'E']
After:
pretrained_dict: ['A', 'B', 'C']
model_dict: ['A', 'B', 'C', 'E']
pretrained_dict = ...
model_dict = model.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load the new state dict
model.load_state_dict(model_dict)
https://discuss.pytorch.org/t/how-to-load-part-of-pre-trained-model/1113/16
本文内容由网友自发贡献，版权归原作者所有，本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容，请联系:hwhale#tublm.com(使用前将#替换为@)
