我正在尝试在 Colab 云 GPU 上使用 HuggingFace 变压器库训练 NER 模型,对其进行 pickle 并将模型加载到我自己的 CPU 上以进行预测。
Code
模型如下:
from transformers import BertForTokenClassification
model = BertForTokenClassification.from_pretrained(
"bert-base-cased",
num_labels=NUM_LABELS,
output_attentions = False,
output_hidden_states = False
)
我正在使用此代码片段将模型保存在 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()
用于标记文本并进行实际预测的代码片段在 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]
有谁知道为什么它不起作用?我错过了什么?