torch.load_state_dict()用法

G_inkk

于 2021-08-26 11:14:59 发布

阅读量4.2k

点赞数 2

文章标签： pytorch 深度学习

本文链接：https://blog.youkuaiyun.com/G_inkk/article/details/119927235

版权

在Pytorch中，使用torchvision.models.resnet50预构建模型，并通过load_state_dict()函数加载预训练权重。当strict参数为True时，模型的层名必须与预训练权重完全匹配，否则会报错。若设置为False，不匹配的权重将被忽略，模型其余部分使用预训练权重。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在Pytorch中构建好一个模型后，一般需要进行预训练权重中加载。torch.load_state_dict()函数就是用于将预训练的参数权重加载到新的模型之中。

net = torchvision.models.resnte50(pretrained=False)
net.load_state_dict(torch.load('*.pth'), strict=True)

当strict=True,要求预训练权重层数的键值与新构建的模型中的权重层数名称完全吻合；如果新构建的模型在层数上进行了部分微调，则上述代码就会报错：说key对应不上。

若strict=False ，训练权重中与新构建网络中匹配层的键值就进行使用，没有的就默认初始化。