加载预训练参数的四种方法

Dream Algorithm

已于 2023-07-30 17:29:47 修改

阅读量1.1k

点赞数 1

CC 4.0 BY-SA版权

文章标签： pytorch 神经网络深度学习 python 人工智能

于 2023-03-10 15:14:52 首次发布

本文链接：https://blog.csdn.net/A2321161581/article/details/129442526

文章讲述了在PyTorch中如何加载预训练模型的权重，包括检查文件存在性，删除特定层的权重，以及自定义模型时从预训练模型中迁移权重的过程。同时，提供了一个resnet18模型预训练和调整最后一层以适应不同类别数目的示例。

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

一、


model_weight_path = ""
assert os.path.exists(model_weight_path), "file {} dose not exist.".format(model_weight_path)
pre_weights = torch.load(model_weight_path, map_location=device)

 # delete classifier weights
pre_dict = {k: v for k, v in pre_weights.items() if net.state_dict()[k].numel() == v.numel()}
net.load_state_dict(pre_dict, strict=False)

二、


weights_dict = torch.load("path",map_location=device)
# 删除有关分类类别的权重
for k in list(weights_dict.keys()):
   if "fc" in k:
      del weights_dict[k]
net.load_state_dict(weights_dict, strict=False)

三、


#加载model，mymodel是自己定义好的模型
    pretrainedmodel = resnet18(pretrained=True) 
    # mymodel =Net(...) 
    
    #读取参数 
    pretrained_dict = pretrainedmodel.state_dict() 
    model_dict = net.state_dict() 
    
    #将pretrained_dict里不属于model_dict的键剔除掉 
    pretrained_dict =  {k: v for k, v in pretrained_dict.items() if k in model_dict} 
    
    # 更新现有的model_dict 
    model_dict.update(pretrained_dict) 
    
    # 加载我们真正需要的state_dict 
    net.load_state_dict(model_dict)

四、


def resnet18(pretrained=False, **kwargs):
    """Constructs a ResNet-18 model.
    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    if kwargs['num_classes'] != 1000 and pretrained:
        model = ResNet(BasicBlock, [2, 2, 2, 2])
        model.load_state_dict(model_zoo.load_url(model_urls['resnet18']))
        model.fc = nn.Linear(model.fc.in_features, kwargs['num_classes'])
        print('done load model')
    else:
        model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)

    return model