LLM 课程文档

词元分类

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

Token 分类

Ask a Question Open In Colab Open In Studio Lab

我们将探索的第一个应用是 token 分类。这项通用任务涵盖了任何可以表述为“为句子中的每个 token 归因一个标签”的问题,例如:

  • 命名实体识别 (NER):在句子中找到实体(如人名、地点或组织)。这可以通过为每个 token 分配一个标签来表述,其中每个实体有一个类别,另外一个类别用于“非实体”。
  • 词性标注 (POS):将句子中的每个词标记为对应的特定词性(如名词、动词、形容词等)。
  • 组块化 (Chunking):找到属于同一实体的 token。这项任务(可以与 POS 或 NER 结合使用)可以表述为:为组块开头的任何 token 分配一个标签(通常是 B-),为组块内部的 token 分配另一个标签(通常是 I-),以及为不属于任何组块的 token 分配第三个标签(通常是 O)。

当然,还有许多其他类型的 token 分类问题;这些只是一些有代表性的例子。在本节中,我们将对一个 NER 任务的模型 (BERT) 进行微调,然后它将能够计算出如下所示的预测:

One-hot encoded labels for question answering.

您可以在此处找到我们将训练并上传到 Hub 的模型,并仔细检查其预测。

准备数据

首先,我们需要一个适合 token 分类的数据集。在本节中,我们将使用CoNLL-2003 数据集,其中包含路透社的新闻报道。

💡 只要你的数据集由分成词语并带有相应标签的文本组成,你就可以将此处描述的数据处理过程应用于你自己的数据集。如果需要复习如何在 Dataset 中加载你自己的自定义数据,请参阅第五章

CoNLL-2003 数据集

要加载 CoNLL-2003 数据集,我们使用 🤗 Datasets 库中的 load_dataset() 方法:

from datasets import load_dataset

raw_datasets = load_dataset("conll2003")

这将下载并缓存数据集,就像我们在第三章中看到的 GLUE MRPC 数据集一样。检查此对象会显示存在的列以及训练集、验证集和测试集之间的划分:

raw_datasets
DatasetDict({
    train: Dataset({
        features: ['chunk_tags', 'id', 'ner_tags', 'pos_tags', 'tokens'],
        num_rows: 14041
    })
    validation: Dataset({
        features: ['chunk_tags', 'id', 'ner_tags', 'pos_tags', 'tokens'],
        num_rows: 3250
    })
    test: Dataset({
        features: ['chunk_tags', 'id', 'ner_tags', 'pos_tags', 'tokens'],
        num_rows: 3453
    })
})

特别地,我们可以看到数据集包含了我们之前提到的三个任务的标签:NER、POS 和组块化。与其他数据集的一个巨大区别是,输入文本不以句子或文档的形式呈现,而是以词语列表的形式(最后一列称为 tokens,但它包含的词语是经过预分词的输入,仍需要通过分词器进行子词分词)。

让我们看看训练集的第一个元素:

raw_datasets["train"][0]["tokens"]
['EU', 'rejects', 'German', 'call', 'to', 'boycott', 'British', 'lamb', '.']

由于我们希望执行命名实体识别,我们将查看 NER 标签:

raw_datasets["train"][0]["ner_tags"]
[3, 0, 7, 0, 0, 0, 7, 0, 0]

这些是准备用于训练的整数标签,但当我们需要检查数据时,它们不一定有用。与文本分类类似,我们可以通过查看数据集的 features 属性来访问这些整数与标签名称之间的对应关系:

ner_feature = raw_datasets["train"].features["ner_tags"]
ner_feature
Sequence(feature=ClassLabel(num_classes=9, names=['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC'], names_file=None, id=None), length=-1, id=None)

所以这一列包含的元素是 ClassLabel 序列。序列元素的类型在 ner_featurefeature 属性中,我们可以通过查看该 featurenames 属性来访问名称列表:

label_names = ner_feature.feature.names
label_names
['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC']

我们在第六章深入了解 token-classification 管道时已经见过这些标签,但为了快速回顾一下:

  • O 表示该词不对应任何实体。
  • B-PER/I-PER 表示该词对应于一个*人物*实体的开头/内部。
  • B-ORG/I-ORG 表示该词对应于一个*组织*实体的开头/内部。
  • B-LOC/I-LOC 表示该词对应于一个*地点*实体的开头/内部。
  • B-MISC/I-MISC 表示该词对应于一个*杂项*实体的开头/内部。

现在,解码我们之前看到的标签会得到:

words = raw_datasets["train"][0]["tokens"]
labels = raw_datasets["train"][0]["ner_tags"]
line1 = ""
line2 = ""
for word, label in zip(words, labels):
    full_label = label_names[label]
    max_length = max(len(word), len(full_label))
    line1 += word + " " * (max_length - len(word) + 1)
    line2 += full_label + " " * (max_length - len(full_label) + 1)

print(line1)
print(line2)
'EU    rejects German call to boycott British lamb .'
'B-ORG O       B-MISC O    O  O       B-MISC  O    O'

对于一个混合了 B-I- 标签的例子,在训练集索引 4 的元素上,相同的代码会给我们:

'Germany \'s representative to the European Union \'s veterinary committee Werner Zwingmann said on Wednesday consumers should buy sheepmeat from countries other than Britain until the scientific advice was clearer .'
'B-LOC   O  O              O  O   B-ORG    I-ORG O  O          O         B-PER  I-PER     O    O  O         O         O      O   O         O    O         O     O    B-LOC   O     O   O          O      O   O       O'

正如我们所看到的,跨越两个词的实体,如“European Union”和“Werner Zwingmann”,第一个词被赋予 B- 标签,第二个词被赋予 I- 标签。

✏️ 轮到你了! 打印带有 POS 或分块标签的相同两个句子。

处理数据

像往常一样,我们的文本需要转换为 token ID,模型才能理解它们。正如我们在第六章中看到的,token 分类任务的一个重要区别是,我们有预分词的输入。幸运的是,tokenizer API 可以很轻松地处理这个问题;我们只需要用一个特殊标志来警告 tokenizer

首先,让我们创建 tokenizer 对象。正如我们之前所说,我们将使用 BERT 预训练模型,所以我们将首先下载并缓存相关的 tokenizer:

from transformers import AutoTokenizer

model_checkpoint = "bert-base-cased"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

你可以用Hub上的任何其他模型替换 model_checkpoint,或者用你保存了预训练模型和 tokenizer 的本地文件夹替换。唯一的限制是 tokenizer 需要由 🤗 Tokenizers 库支持,因此有一个“fast”版本可用。你可以在这张大表中看到所有带有 fast 版本的架构,并且要检查你正在使用的 tokenizer 对象是否确实由 🤗 Tokenizers 支持,你可以查看它的 is_fast 属性:

tokenizer.is_fast
True

要对预分词的输入进行分词,我们可以像往常一样使用我们的 tokenizer,只需添加 is_split_into_words=True

inputs = tokenizer(raw_datasets["train"][0]["tokens"], is_split_into_words=True)
inputs.tokens()
['[CLS]', 'EU', 'rejects', 'German', 'call', 'to', 'boycott', 'British', 'la', '##mb', '.', '[SEP]']

正如我们所看到的,分词器添加了模型使用的特殊标记(开头是 [CLS],结尾是 [SEP]),并且大部分单词保持不变。然而,单词 lamb 被分词为两个子词:la##mb。这导致了我们的输入和标签之间存在不匹配:标签列表只有 9 个元素,而我们的输入现在有 12 个标记。处理特殊标记很容易(我们知道它们在开头和结尾),但我们还需要确保将所有标签与正确的单词对齐。

幸运的是,由于我们使用的是快速分词器,我们可以使用 🤗 Tokenizers 的超能力,这意味着我们可以轻松地将每个 token 映射到其对应的单词(如第六章中所述):

inputs.word_ids()
[None, 0, 1, 2, 3, 4, 5, 6, 7, 7, 8, None]

稍作修改,我们就可以扩展标签列表以匹配 token。我们将应用的第一条规则是特殊 token 的标签为 -100。这是因为默认情况下,-100 是我们稍后将使用的损失函数(交叉熵)中忽略的索引。然后,每个 token 都获得与其所在单词开头的 token 相同的标签,因为它们属于同一个实体。对于单词内部但不在开头的 token,我们将 B- 替换为 I-(因为该 token 不以实体开头)。

def align_labels_with_tokens(labels, word_ids):
    new_labels = []
    current_word = None
    for word_id in word_ids:
        if word_id != current_word:
            # Start of a new word!
            current_word = word_id
            label = -100 if word_id is None else labels[word_id]
            new_labels.append(label)
        elif word_id is None:
            # Special token
            new_labels.append(-100)
        else:
            # Same word as previous token
            label = labels[word_id]
            # If the label is B-XXX we change it to I-XXX
            if label % 2 == 1:
                label += 1
            new_labels.append(label)

    return new_labels

让我们在我们的第一个句子上试试看:

labels = raw_datasets["train"][0]["ner_tags"]
word_ids = inputs.word_ids()
print(labels)
print(align_labels_with_tokens(labels, word_ids))
[3, 0, 7, 0, 0, 0, 7, 0, 0]
[-100, 3, 0, 7, 0, 0, 0, 7, 0, 0, 0, -100]

正如我们所看到的,我们的函数在开头和结尾添加了两个特殊 token 的 -100,并为我们被拆分为两个 token 的单词添加了一个新的 0

✏️ 轮到你了! 有些研究人员更喜欢每个词只分配一个标签,并将 -100 分配给给定词中的其他子 token。这是为了避免那些被拆分成许多子 token 的长词对损失函数造成过大的影响。请按照此规则更改前面的函数,以将标签与输入 ID 对齐。

为了预处理我们的整个数据集,我们需要对所有输入进行分词,并对所有标签应用 align_labels_with_tokens()。为了利用我们快速分词器的速度,最好同时对大量文本进行分词,所以我们将编写一个函数来处理示例列表,并使用带有 batched=True 选项的 Dataset.map() 方法。与我们之前的示例唯一不同的是,当分词器的输入是文本列表(或者在我们的例子中是词语列表的列表)时,word_ids() 函数需要获取我们想要其词语 ID 的示例的索引,所以我们也添加了这一点:

def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(
        examples["tokens"], truncation=True, is_split_into_words=True
    )
    all_labels = examples["ner_tags"]
    new_labels = []
    for i, labels in enumerate(all_labels):
        word_ids = tokenized_inputs.word_ids(i)
        new_labels.append(align_labels_with_tokens(labels, word_ids))

    tokenized_inputs["labels"] = new_labels
    return tokenized_inputs

请注意,我们尚未填充输入;我们将在稍后使用数据收集器创建批次时进行填充。

现在我们可以一次性地对数据集的其他拆分应用所有这些预处理:

tokenized_datasets = raw_datasets.map(
    tokenize_and_align_labels,
    batched=True,
    remove_columns=raw_datasets["train"].column_names,
)

最困难的部分已经完成了!现在数据已经预处理完毕,实际的训练将与我们在第三章中所做的非常相似。

使用 Trainer API 微调模型

使用 Trainer 的实际代码将与之前相同;唯一的更改是数据如何整理成批次以及度量计算函数。

数据整理

我们不能像第三章那样简单地使用 DataCollatorWithPadding,因为它只填充输入(输入 ID、注意力掩码和 token 类型 ID)。在这里,我们的标签应该以与输入完全相同的方式进行填充,以便它们保持相同的大小,使用 -100 作为值,这样相应的预测在损失计算中就会被忽略。

这一切都由一个DataCollatorForTokenClassification完成。与 DataCollatorWithPadding 一样,它也接受用于预处理输入的 tokenizer

from transformers import DataCollatorForTokenClassification

data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)

要在几个样本上进行测试,我们只需在来自我们 token 化的训练集的一些示例列表上调用它:

batch = data_collator([tokenized_datasets["train"][i] for i in range(2)])
batch["labels"]
tensor([[-100,    3,    0,    7,    0,    0,    0,    7,    0,    0,    0, -100],
        [-100,    1,    2, -100, -100, -100, -100, -100, -100, -100, -100, -100]])

让我们将它与数据集中第一个和第二个元素的标签进行比较:

for i in range(2):
    print(tokenized_datasets["train"][i]["labels"])
[-100, 3, 0, 7, 0, 0, 0, 7, 0, 0, 0, -100]
[-100, 1, 2, -100]

正如我们所看到的,第二组标签已被填充到第一组的长度,使用了 -100

指标

为了让 Trainer 在每个 epoch 计算一个度量,我们需要定义一个 compute_metrics() 函数,该函数接受预测和标签数组,并返回一个包含度量名称和值的字典。

用于评估 token 分类预测的传统框架是seqeval。要使用此度量,我们首先需要安装 seqeval 库:

!pip install seqeval

然后我们可以像在第三章中那样通过 evaluate.load() 函数加载它:

import evaluate

metric = evaluate.load("seqeval")

这个指标不像标准的准确率那样工作:它实际上会将标签列表作为字符串而不是整数,所以我们需要在将它们传递给指标之前完全解码预测和标签。让我们看看它是如何工作的。首先,我们将获取第一个训练示例的标签:

labels = raw_datasets["train"][0]["ner_tags"]
labels = [label_names[i] for i in labels]
labels
['B-ORG', 'O', 'B-MISC', 'O', 'O', 'O', 'B-MISC', 'O', 'O']

然后我们可以通过仅更改索引 2 处的值来创建它们的假预测:

predictions = labels.copy()
predictions[2] = "O"
metric.compute(predictions=[predictions], references=[labels])

请注意,此度量接受预测列表(不仅仅是一个)和标签列表。以下是输出:

{'MISC': {'precision': 1.0, 'recall': 0.5, 'f1': 0.67, 'number': 2},
 'ORG': {'precision': 1.0, 'recall': 1.0, 'f1': 1.0, 'number': 1},
 'overall_precision': 1.0,
 'overall_recall': 0.67,
 'overall_f1': 0.8,
 'overall_accuracy': 0.89}

这会返回大量信息!我们获得了每个独立实体的精确率、召回率和 F1 分数,以及总体分数。对于我们的指标计算,我们将只保留总体分数,但您可以随意调整 compute_metrics() 函数以返回您希望报告的所有指标。

这个 compute_metrics() 函数首先对 logits 进行 argmax 以将其转换为预测(通常,logits 和概率的顺序相同,因此我们不需要应用 softmax)。然后我们必须将标签和预测都从整数转换为字符串。我们删除所有标签为 -100 的值,然后将结果传递给 metric.compute() 方法:

import numpy as np


def compute_metrics(eval_preds):
    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)

    # Remove ignored index (special tokens) and convert to labels
    true_labels = [[label_names[l] for l in label if l != -100] for label in labels]
    true_predictions = [
        [label_names[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    all_metrics = metric.compute(predictions=true_predictions, references=true_labels)
    return {
        "precision": all_metrics["overall_precision"],
        "recall": all_metrics["overall_recall"],
        "f1": all_metrics["overall_f1"],
        "accuracy": all_metrics["overall_accuracy"],
    }

现在,这一切都完成了,我们几乎可以定义我们的 Trainer 了。我们只需要一个要微调的 model

定义模型

由于我们正在处理 token 分类问题,我们将使用 AutoModelForTokenClassification 类。定义此模型时要记住的主要事情是传递我们拥有的标签数量的一些信息。最简单的方法是使用 num_labels 参数传递该数量,但如果我们想要一个像本节开头所示那样工作的良好推理小部件,最好设置正确的标签对应关系。

它们应该由两个字典设置:id2labellabel2id,它们包含从 ID 到标签以及从标签到 ID 的映射:

id2label = {i: label for i, label in enumerate(label_names)}
label2id = {v: k for k, v in id2label.items()}

现在我们可以直接将它们传递给 AutoModelForTokenClassification.from_pretrained() 方法,它们将被设置在模型的配置中,然后正确地保存并上传到 Hub:

from transformers import AutoModelForTokenClassification

model = AutoModelForTokenClassification.from_pretrained(
    model_checkpoint,
    id2label=id2label,
    label2id=label2id,
)

就像我们在第三章中定义 AutoModelForSequenceClassification 时一样,创建模型会发出警告,指出某些权重(来自预训练头部)未被使用,而另一些权重(来自新的 token 分类头部)是随机初始化的,并且该模型应该进行训练。我们将在稍后进行训练,但首先让我们仔细检查我们的模型是否具有正确的标签数量:

model.config.num_labels
9

⚠️ 如果您的模型标签数量不正确,稍后调用 Trainer.train() 方法时会遇到一个模糊的错误(例如“CUDA error: device-side assert triggered”)。这是用户报告此类错误的头号原因,因此请务必进行此检查以确认您拥有预期的标签数量。

微调模型

我们现在准备训练我们的模型了!在定义 Trainer 之前,我们只需要做最后两件事:登录 Hugging Face 并定义我们的训练参数。如果你在笔记本中工作,有一个便捷函数可以帮助你完成此操作:

from huggingface_hub import notebook_login

notebook_login()

这将显示一个可以输入您的 Hugging Face 登录凭据的小部件。

如果您不在笔记本中工作,只需在终端中输入以下行

huggingface-cli login

完成此操作后,我们可以定义 TrainingArguments

from transformers import TrainingArguments

args = TrainingArguments(
    "bert-finetuned-ner",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    num_train_epochs=3,
    weight_decay=0.01,
    push_to_hub=True,
)

您以前见过其中的大部分:我们设置了一些超参数(如学习率、训练 epoch 数量和权重衰减),并指定 push_to_hub=True 以表明我们希望保存模型并在每个 epoch 结束时对其进行评估,并且我们希望将结果上传到模型 Hub。请注意,您可以使用 hub_model_id 参数指定要推送到的存储库名称(特别是,您将不得不使用此参数推送到组织)。例如,当我们将模型推送到huggingface-course 组织时,我们向 TrainingArguments 添加了 hub_model_id="huggingface-course/bert-finetuned-ner"。默认情况下,使用的存储库将在您的命名空间中,并以您设置的输出目录命名,因此在我们的例子中将是 "sgugger/bert-finetuned-ner"

💡 如果您使用的输出目录已经存在,它需要是您要推送到的存储库的本地克隆。如果不是,您在定义 Trainer 时会收到错误,并且需要设置一个新名称。

最后,我们只需将所有内容传递给 Trainer 并启动训练:

from transformers import Trainer

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    processing_class=tokenizer,
)
trainer.train()

请注意,在训练过程中,每次模型保存时(此处是每个 epoch),它都会在后台上传到 Hub。这样,如有必要,您就可以在另一台机器上恢复训练。

训练完成后,我们使用 push_to_hub() 方法确保上传模型的最新版本:

trainer.push_to_hub(commit_message="Training complete")

此命令返回刚刚提交的 URL,如果您想检查它:

'https://huggingface.co/sgugger/bert-finetuned-ner/commit/26ab21e5b1568f9afeccdaed2d8715f571d786ed'

训练器还会起草一份包含所有评估结果的模型卡并将其上传。在此阶段,您可以使用模型中心上的推理小部件来测试您的模型并与您的朋友分享。您已成功微调了 token 分类任务模型——恭喜!

如果您想更深入地了解训练循环,我们现在将向您展示如何使用 🤗 Accelerate 来完成同样的事情。

自定义训练循环

现在让我们来看看完整的训练循环,这样您就可以轻松定制您需要的部分。它将与我们在第三章中所做的非常相似,只是评估部分有一些变化。

为训练做准备

首先,我们需要从数据集中构建 DataLoader。我们将重用 data_collator 作为 collate_fn,并打乱训练集,但不打乱验证集:

from torch.utils.data import DataLoader

train_dataloader = DataLoader(
    tokenized_datasets["train"],
    shuffle=True,
    collate_fn=data_collator,
    batch_size=8,
)
eval_dataloader = DataLoader(
    tokenized_datasets["validation"], collate_fn=data_collator, batch_size=8
)

接下来我们重新实例化模型,以确保我们不是从之前继续微调,而是从 BERT 预训练模型重新开始:

model = AutoModelForTokenClassification.from_pretrained(
    model_checkpoint,
    id2label=id2label,
    label2id=label2id,
)

然后我们需要一个优化器。我们将使用经典的 AdamW,它类似于 Adam,但修复了权重衰减的应用方式:

from torch.optim import AdamW

optimizer = AdamW(model.parameters(), lr=2e-5)

一旦我们拥有了所有这些对象,我们就可以将它们发送到 accelerator.prepare() 方法:

from accelerate import Accelerator

accelerator = Accelerator()
model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(
    model, optimizer, train_dataloader, eval_dataloader
)

🚨 如果您正在 TPU 上训练,您需要将从上面单元格开始的所有代码移动到一个专门的训练函数中。有关更多详细信息,请参阅第三章

现在我们已经将 train_dataloader 发送给了 accelerator.prepare(),我们可以使用它的长度来计算训练步骤的数量。请记住,我们应该总是在准备完数据加载器之后再执行此操作,因为该方法会更改其长度。我们使用经典的从学习率到 0 的线性调度:

from transformers import get_scheduler

num_train_epochs = 3
num_update_steps_per_epoch = len(train_dataloader)
num_training_steps = num_train_epochs * num_update_steps_per_epoch

lr_scheduler = get_scheduler(
    "linear",
    optimizer=optimizer,
    num_warmup_steps=0,
    num_training_steps=num_training_steps,
)

最后,为了将我们的模型推送到 Hub,我们需要在一个工作文件夹中创建一个 Repository 对象。如果尚未登录,请先登录 Hugging Face。我们将根据我们希望给模型起的模型 ID 来确定存储库名称(请随意用您自己的选择替换 repo_name;它只需要包含您的用户名,这就是 get_full_repo_name() 函数所做的工作):

from huggingface_hub import Repository, get_full_repo_name

model_name = "bert-finetuned-ner-accelerate"
repo_name = get_full_repo_name(model_name)
repo_name
'sgugger/bert-finetuned-ner-accelerate'

然后我们可以在本地文件夹中克隆该存储库。如果它已经存在,则该本地文件夹应该是我们正在使用的存储库的现有克隆:

output_dir = "bert-finetuned-ner-accelerate"
repo = Repository(output_dir, clone_from=repo_name)

现在我们可以通过调用 repo.push_to_hub() 方法上传我们保存在 output_dir 中的任何内容。这将帮助我们在每个 epoch 结束时上传中间模型。

训练循环

我们现在准备编写完整的训练循环。为了简化其评估部分,我们定义了 postprocess() 函数,该函数接受预测和标签,并将其转换为字符串列表,正如我们的 metric 对象所期望的那样:

def postprocess(predictions, labels):
    predictions = predictions.detach().cpu().clone().numpy()
    labels = labels.detach().cpu().clone().numpy()

    # Remove ignored index (special tokens) and convert to labels
    true_labels = [[label_names[l] for l in label if l != -100] for label in labels]
    true_predictions = [
        [label_names[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    return true_labels, true_predictions

然后我们可以编写完整的训练循环。在定义进度条以跟踪训练进度之后,循环包含三个部分:

  • 训练本身,这是对 train_dataloader 的经典迭代,模型前向传播,然后反向传播和优化器步进。
  • 评估部分,在获取模型在批次上的输出后,有一个新颖之处:由于两个进程可能将输入和标签填充成不同的形状,我们需要使用 accelerator.pad_across_processes() 来使预测和标签具有相同的形状,然后才能调用 gather() 方法。如果我们不这样做,评估将要么出错,要么永远挂起。然后我们将结果发送到 metric.add_batch(),并在评估循环结束后调用 metric.compute()
  • 保存和上传,我们首先保存模型和 tokenizer,然后调用 repo.push_to_hub()。请注意,我们使用参数 blocking=False 来告诉 🤗 Hub 库以异步进程推送。这样,训练正常进行,并且此(长时间)指令在后台执行。

这是训练循环的完整代码:

from tqdm.auto import tqdm
import torch

progress_bar = tqdm(range(num_training_steps))

for epoch in range(num_train_epochs):
    # Training
    model.train()
    for batch in train_dataloader:
        outputs = model(**batch)
        loss = outputs.loss
        accelerator.backward(loss)

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)

    # Evaluation
    model.eval()
    for batch in eval_dataloader:
        with torch.no_grad():
            outputs = model(**batch)

        predictions = outputs.logits.argmax(dim=-1)
        labels = batch["labels"]

        # Necessary to pad predictions and labels for being gathered
        predictions = accelerator.pad_across_processes(predictions, dim=1, pad_index=-100)
        labels = accelerator.pad_across_processes(labels, dim=1, pad_index=-100)

        predictions_gathered = accelerator.gather(predictions)
        labels_gathered = accelerator.gather(labels)

        true_predictions, true_labels = postprocess(predictions_gathered, labels_gathered)
        metric.add_batch(predictions=true_predictions, references=true_labels)

    results = metric.compute()
    print(
        f"epoch {epoch}:",
        {
            key: results[f"overall_{key}"]
            for key in ["precision", "recall", "f1", "accuracy"]
        },
    )

    # Save and upload
    accelerator.wait_for_everyone()
    unwrapped_model = accelerator.unwrap_model(model)
    unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)
    if accelerator.is_main_process:
        tokenizer.save_pretrained(output_dir)
        repo.push_to_hub(
            commit_message=f"Training in progress epoch {epoch}", blocking=False
        )

如果这是您第一次看到使用 🤗 Accelerate 保存的模型,让我们花点时间检查一下相关的三行代码:

accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save)

第一行不言自明:它告诉所有进程等待,直到所有进程都到达该阶段才继续。这是为了确保我们在保存之前所有进程中的模型都相同。然后我们获取 unwrapped_model,这是我们定义的原始模型。accelerator.prepare() 方法会更改模型以在分布式训练中工作,因此它将不再具有 save_pretrained() 方法;accelerator.unwrap_model() 方法会撤销该步骤。最后,我们调用 save_pretrained() 但告诉该方法使用 accelerator.save() 而不是 torch.save()

完成此操作后,您应该会得到一个与使用 Trainer 训练的模型产生非常相似结果的模型。您可以在huggingface-course/bert-finetuned-ner-accelerate查看我们使用此代码训练的模型。如果您想测试对训练循环的任何调整,可以直接通过编辑上面显示的代码来实现!

使用微调模型

我们已经向您展示了如何在模型 Hub 上使用我们微调的模型进行推理。要在 pipeline 中本地使用它,您只需指定正确的模型标识符:

from transformers import pipeline

# Replace this with your own checkpoint
model_checkpoint = "huggingface-course/bert-finetuned-ner"
token_classifier = pipeline(
    "token-classification", model=model_checkpoint, aggregation_strategy="simple"
)
token_classifier("My name is Sylvain and I work at Hugging Face in Brooklyn.")
[{'entity_group': 'PER', 'score': 0.9988506, 'word': 'Sylvain', 'start': 11, 'end': 18},
 {'entity_group': 'ORG', 'score': 0.9647625, 'word': 'Hugging Face', 'start': 33, 'end': 45},
 {'entity_group': 'LOC', 'score': 0.9986118, 'word': 'Brooklyn', 'start': 49, 'end': 57}]

太棒了!我们的模型与此管道的默认模型一样出色!

< > 在 GitHub 上更新