pytorch pretrained model

最新推荐文章于 2021-06-04 10:14:53 发布

原创最新推荐文章于 2021-06-04 10:14:53 发布 · 3.1k 阅读

1 ·

CC 4.0 BY-SA版权

无专栏收录该内容

6 篇文章

订阅专栏

本文介绍两种在PyTorch中加载预训练模型的方法：一种是直接比较模型大小进行加载，另一种是通过筛选预训练权重中的键来匹配模型状态字典。

部署运行你感兴趣的模型镜像

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

您可能感兴趣的与本文相关的镜像