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

最新推荐文章于 2024-08-08 22:46:59 发布

原创最新推荐文章于 2024-08-08 22:46:59 发布 · 652 阅读

0 ·

CC 4.0 BY-SA版权

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

2 篇文章

订阅专栏

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

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

    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])

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