小编fly*_*ans的帖子

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

代码

模型如下：

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 …