首页 > 解决方案 > RuntimeError:为 BertModel 加载 state_dict 时出错

问题描述

我使用拥抱人脸转换器库对 BERT 模型进行微调,并在云端的 GPU 中对其进行训练。然后我保存模型和标记器,如下所示:

model.save_pretrained('/saved_model/')
torch.save(best_model.state_dict(), '/saved_model/model')
tokenizer.save_pretrained('/saved_model/')

我将saved_model目录下载到我的计算机中。然后我在我的计算机中加载如下所示的模型/标记器

import torch
from transformers import *
tokenizer = BertTokenizer.from_pretrained('./saved_model/')
config = BertConfig('./saved_model/config.json')
model = BertModel(config)
model.load_state_dict(torch.load('./saved_model/pytorch_model.bin', map_location=torch.device('cpu')))
model.eval()

但它会为该model.load_state_dict行抛出错误

RuntimeError: Error(s) in loading state_dict for BertModel:
    Missing key(s) in state_dict:

它列出了 state_dict 中显然缺少的一堆键。

我是 pytorch 的新手,不知道发生了什么。很可能我没有以正确的方式保存模型。

请建议。

标签: machine-learningdeep-learningpytorchtransfer-learning

解决方案


您可能知道,state_dictPyTorch 模块的OrderedDict. 当您尝试从 a 加载模块的权重时state_dict,它会抱怨缺少键,这意味着state_dict不包含这些键。在这种情况下,我建议采取以下措施。

  1. 检查哪些键存在于state_dict. 听起来不可能只保存密钥的子集。
  2. 此外,请确保您已加载正确的配置。否则,如果您训练的 BertModel 和您要为其加载权重的新 BertModel 不同,那么您将收到此错误。
  3. 最后,如果您的代码通过上述两种情况,然后保存模型,请确保将所有层的参数保存在文件中。该声明torch.save(best_model.state_dict(), '/saved_model/model')对我来说看起来不错,但请确保best_model.state_dict()包含所有预期的键。

推荐阅读