PyTorch加载已保存模型时出现KeyError问题求助

核心问题分析 #

KeyError本质是加载后的模型遇到了词汇表中未收录的token，你怀疑的PAD token差异是核心排查方向，同时还要聚焦词汇表一致性、模型保存/加载配置完整性这两个关键因素。

具体排查与解决步骤 #

1. 严格对齐Tokenizer（词汇表） #

• 训练时使用的Tokenizer必须和加载模型时完全一致，禁止重新初始化或使用不同的词汇表文件。
  • 解决方式：训练完成后务必保存Tokenizer（以Hugging Face框架为例）：

    tokenizer.save_pretrained("saved_bengali_tokenizer/")
    

    加载模型时直接复用该Tokenizer：

    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained("saved_bengali_tokenizer/")
    

• 检查加载时的自动padding配置：部分Tokenizer默认会在批量处理时自动添加PAD，若训练时是手动控制padding逻辑，需显式统一策略：

  # 训练和加载时保持相同的padding设置
  tokenizer = AutoTokenizer.from_pretrained("saved_bengali_tokenizer/", padding_side="right", padding=False)
  

2. 确保模型保存与加载的完整性 #

• 避免手动保存模型权重（如torch.save），这类方式会遗漏模型配置（如词汇表大小、PAD token ID等），导致embedding层与Tokenizer不匹配。
  • 解决方式：使用框架原生的完整保存方法，以Hugging Face为例：

    model.save_pretrained("saved_bengali_stemmer/")
    

    加载时直接读取完整模型：

    from transformers import AutoModelForTokenClassification
    model = AutoModelForTokenClassification.from_pretrained("saved_bengali_stemmer/")
    

• 加载后确认pad_token_id配置：训练时若设置过model.config.pad_token_id = tokenizer.pad_token_id，加载后需重新验证该参数是否被正确继承，避免模型无法识别PAD token。

3. 对齐输入处理逻辑 #

• 对比同一运行时与新运行时的输入处理代码，确保tokenization流程完全一致：
  • 检查是否在新运行时额外添加了padding、截断操作，或训练时用了特殊的文本预处理逻辑（如孟加拉语的词分割规则），加载时未同步。
  • 示例：训练和加载时统一输入处理参数：

    # 两地代码保持完全一致
    inputs = tokenizer(texts, truncation=True, padding="max_length", max_length=64, return_tensors="pt")
    

4. 孟加拉语Tokenizer特殊处理 #

• 孟加拉语属于低资源语言，部分预训练Tokenizer可能未默认设置PAD token，需手动添加并同步模型embedding层：
  • 训练前执行：

    if tokenizer.pad_token is None:
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
        model.resize_token_embeddings(len(tokenizer))
    

  • 加载模型时，确保Tokenizer的PAD token与训练时完全一致，避免出现未收录的token。

内容的提问来源于stack exchange，提问作者Mahir RAfid