基于 BERT 的 NER 模型在反序列化时给出不一致的预测

Posted 2023-03-12

【中文标题】基于 BERT 的 NER 模型在反序列化时给出不一致的预测

【英文标题】：BERT-based NER model giving inconsistent prediction when deserialized

【发布时间】：2021-02-13 01:28:17

【问题描述】：

我正在尝试在 Colab 云 GPU 上使用 HuggingFace 转换器库训练 NER 模型，将其腌制并在我自己的 CPU 上加载模型以进行预测。

代码

型号如下：

from transformers import BertForTokenClassification

model = BertForTokenClassification.from_pretrained(
    "bert-base-cased",
    num_labels=NUM_LABELS,
    output_attentions = False,
    output_hidden_states = False
)

我正在使用这个 sn-p 将模型保存在 Colab 上

import torch

torch.save(model.state_dict(), FILENAME)

然后将其加载到我的本地 CPU 上使用

# Initiating an instance of the model type

model_reload = BertForTokenClassification.from_pretrained(
    "bert-base-cased",
    num_labels=len(tag2idx),
    output_attentions = False,
    output_hidden_states = False
)

# Loading the model
model_reload.load_state_dict(torch.load(FILENAME, map_location='cpu'))
model_reload.eval()

用于标记文本和进行实际预测的代码 sn-p 在 Colab GPU 笔记本实例和我的 CPU 笔记本实例上都是相同的。

预期行为

经过 GPU 训练的模型表现正确并完美分类以下标记：

O       [CLS]
O       Good
O       morning
O       ,
O       my
O       name
O       is
B-per   John
I-per   Kennedy
O       and
O       I
O       am
O       working
O       at
B-org   Apple
O       in
O       the
O       headquarters
O       of
B-geo   Cupertino
O       [SEP]

实际行为

当加载模型并使用它在我的 CPU 上进行预测时，预测完全错误：

I-eve   [CLS]
I-eve   Good
I-eve   morning
I-eve   ,
I-eve   my
I-eve   name
I-eve   is
I-geo   John
B-eve   Kennedy
I-eve   and
I-eve   I
I-eve   am
I-eve   working
I-eve   at
I-gpe   Apple
I-eve   in
I-eve   the
I-eve   headquarters
I-eve   of
B-org   Cupertino
I-eve   [SEP]

有没有人知道为什么它不起作用？我错过了什么吗？

【问题讨论】：

你能和我们分享你的state_dict吗？

还有两个关于保存和加载变压器模型的问题会导致精度降低（here 和 here）。我认为你应该打开一个错误报告。

【参考方案1】：

我修好了，有两个问题：

标记的索引标签映射错误，由于某种原因，list() 函数在 Colab GPU 上的工作方式与我的 CPU 不同 (??)

用于保存模型的sn-p不正确，对于基于huggingface-transformers库的模型你不能使用model.save_dict()并稍后加载它，你需要使用save_pretrained()方法您的模型类，稍后使用 from_pretrained() 加载它。