工具代码粘贴3——pytorch载入模型

Mona-abc

于 2019-07-21 14:54:37 发布

阅读量621

点赞数

CC 4.0 BY-SA版权

分类专栏： pytorch代码粘贴

本文链接：https://blog.youkuaiyun.com/csnc007/article/details/96722282

pytorch代码粘贴专栏收录该内容

2 篇文章

订阅专栏

本文介绍了两种加载预训练模型的方法。第一种方法使用PyTorch的state_dict进行参数加载，通过匹配模型字典和预训练权重来实现。第二种方法通过定义load_param函数，逐层复制参数，对于全连接层(fc)则跳过加载。

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

    print('loading checkpoint.......')
    model_dict = model.state_dict()
    pretrained_dict = torch.load(weight_path)
    pretrained_dict = {k:v for k, v in pretrained_dict.items() if k in model_dict}
    model_dict.update(pretrained_dict)
    model.load_state_dict(model_dict)
    print('loaded successfully!')

另：

 def load_param(self, model_path):
        param_dict = torch.load(model_path)
        for i in param_dict:
            if 'fc' in i:
                continue
            self.state_dict()[i].copy_(param_dict[i])