导入模型权重并过滤

angrident

于 2024-07-22 16:02:57 发布

阅读量444

点赞数 5

文章标签： python pytorch

本文链接：https://blog.csdn.net/angrident/article/details/140612446

版权

            
        model_path = './mprnet-jdrl-sdd.pth'  
        old_state_dict = torch.load(model_path)['state_dict']
        # 移除 'module.' 前缀
        old_state_dict = {k.replace('module.', ''): v for k, v in old_state_dict.items()}
        # 过滤掉state_dict中不在新模型中的键
        new_state_dict = {k: v for k, v in old_state_dict.items() if k in self.deblur_model.state_dict()}
        self.deblur_model.load_state_dict(new_state_dict, strict=False)