pytorch使用torch.load加载模型出现错误Unexpected key(s) in state_dict

最新推荐文章于 2024-04-27 20:58:20 发布

夕阳与影

最新推荐文章于 2024-04-27 20:58:20 发布

阅读量3.5k

点赞数 3

分类专栏： deep learning 文章标签：深度学习 pytorch

本文链接：https://blog.youkuaiyun.com/qq_41648043/article/details/108412157

版权

在使用PyTorch的torch.load加载模型时遇到错误，模型状态字典中的key前缀有'module.'，导致加载不匹配。尝试调整代码顺序并检查是否在训练时使用了nn.DataParallel，确保测试时也使用相同设置，问题得到解决。

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

记录错误积累经验
在这里插入图片描述
猛然来这么一大串错误信息让我有点震惊。
pytorch模型文件以key，value形式保存，现在模型所有key都对不上，有点奇怪。细看发现加载的模型相较于当前模型其key值前面多了“module."。
初步诊断

class Tester(object):

    TestParams = TestParams

    def __init__(self, model, test_params):
        assert isinstance(test_params, TestParams)
        self.params = test_params

        # load model
        self.model = model

        ckpt = self.params.ckpt
        if ckpt is not None:
            self.load_ckpt(ckpt)
            logger.info('Load ckpt from {}'.format(ckpt))

        # set CUDA_VISIBLE_DEVICES, 1 GPU is enough
        if len(self.params.gpus) > 0:
            gpu_test = str(self.params.gpus[0])
            os