您的位置:首页 > 其它

pytorch 模型部分参数的加载

2019-02-20 15:47 387 查看

如果对预训练模型的结构进行了一些改动,在训练的开始前希望加载未改动部分的参数,如将resnet18的第一层卷积层conv1的输入由3通道改为6通道的new_conv1,将分类层fc的1000类输出改为2类输出的new_fc,注意:要改一下名字与原来的不同。

导入模型

myNet=ResNet()
然后就加载模型的参数,参考pytorch 如何加载部分预训练模型

pretrained_dict=torch.load(model_weight)
model_dict=myNet.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)
myNet.load_state_dict(model_dict)
也可以通过pretrained model.state_dict()提取需要的模型参数。

作者:lxx516
来源:CSDN
原文:https://blog.csdn.net/LXX516/article/details/80124768
版权声明:本文为博主原创文章,转载请附上博文链接!

内容来自用户分享和网络整理,不保证内容的准确性,如有侵权内容,可联系管理员处理 点击这里给我发消息
标签: