携手创作，共同成长！这是我参与「掘金日新计划 · 8 月更文挑战」的第30天，简单的文本分类任务：使用Huggingface Trainer实现 为大家展示了如何使用Huggingface Transformers库来快速上手实现一个现代的文本分类模型的范式，虽然trainer很好用，但由于其封装的太过完善，所以在某些修改方面显得不是那么方便，我们有必要自己使用Pytorch手写一下整个流程。

任务简介

与前一篇博客相同，仍进行IMDB数据集上的情感分析。不过本次实验的不同之处在于：

1. 为了快速实现，使用训练数据为1000条，测试为100条；
2. 使用BERT进行训练而不是distil-bert。

数据预处理

一如既往，我们还是使用datasets库加载数据集，

from datasets import load_dataset

dataset = load_dataset("imdb")

接着，初始化tokenizer然后进行预处理

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

我们使用bert的tokenizer，词表数目为30522。

在进行分词时，为了加快运行速度，设置最大的token序列长度为256。

def tokenize_function(examples):
    return tokenizer(examples["text"], padding="max_length",truncation=True,max_length=256)

tokenized_datasets = dataset.map(tokenize_function, batched=True)
tokenized_datasets.set_format("torch")  #将值转换为torch.Tensor对象

同样的，由于数据集过大，我们这里只选择1000个做train的样例，100个做evaluation样例。

SMALL_TRAIN_SIZE = 1000
SMALL_TEST_SIZE = 100

small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(SMALL_TRAIN_SIZE))
small_test_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(SMALL_TEST_SIZE))

然后我们使用Pytorch原生的Dataloader来进行数据迭代，这里设置batch size为16，对训练和测试数据集采取相同的batch size：

train_dataloader = DataLoader(small_train_dataset, shuffle=True, batch_size=16)
test_dataloader = DataLoader(small_test_dataset, batch_size=16)

评价指标定义

这里，我们仅选择准确率作为评价指标，

metric=datasets.load_metric("accuracy")

而datasets库已经为我们封装好了评价指标，可以直接一行代码调用，非常方便。

模型、优化器和scheduler的初始化

我们直接使用Huggingface的bert模型进行初始化，同时选择AdamW作为优化器，scheduler则是指的学习率的调度器，用于控制学习率的变化方式，比如BERT中常用的warmup操作。

model=AutoModelForSequenceClassification.from_pretrained("bert-base-uncased",num_labels=2)

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

num_epochs = 1
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
    name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps
)

这里，我们设置训练的epoch为1，学习率为2e-5，batch size已经在Dataloader中设置过为16.

模型的训练与验证

接下来，我们使用Pytorch手写训练和验证流程。

progress_bar = tqdm(range(num_training_steps))
global_step = 0

print("Before training.")
metric=datasets.load_metric("accuracy")
model.eval()
for batch in test_dataloader:
    batch = {k: v.to(device) for k, v in batch.items()}
    with torch.no_grad():
        outputs = model(**batch)

    logits = outputs.logits
    predictions = torch.argmax(logits, dim=-1)
    metric.add_batch(predictions=predictions, references=batch["labels"])

print("step: ", global_step, metric.compute())


print()
print("Start training.")

model.train()
for epoch in range(num_epochs):
    for batch in train_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        loss.backward()

        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
        progress_bar.update(1)
        
        if (global_step+1) % 10 == 0:
            metric=datasets.load_metric("accuracy")
            model.eval()
            for batch in test_dataloader:
                batch = {k: v.to(device) for k, v in batch.items()}
                with torch.no_grad():
                    outputs = model(**batch)

                logits = outputs.logits
                predictions = torch.argmax(logits, dim=-1)
                metric.add_batch(predictions=predictions, references=batch["labels"])
            
            print("step: ", global_step+1, metric.compute())
        global_step += 1
        
    metric=datasets.load_metric("accuracy")
    model.eval()
    for batch in test_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        with torch.no_grad():
            outputs = model(**batch)

        logits = outputs.logits
        predictions = torch.argmax(logits, dim=-1)
        metric.add_batch(predictions=predictions, references=batch["labels"])
    print()
    print("Finish: ", global_step+1, metric.compute())

整体的代码分为三个部分：

1. 训练前的验证：用以检查模型在训练之前的性能表现。
2. 模型训练：包括以下流程，
   • 从Dataloader中读取数据
   • 输入给模型进行前向计算
   • 计算loss数值
   • loss进行梯度回传
   • 优化器进行参数更新
   • 优化器梯度清零
3. 训练完成后进行验证：检验模型的最终表现

得到的打印输出如下：

Before training.
step:  0 {'accuracy': 0.47}

Start training.
step:  10 {'accuracy': 0.53}
step:  20 {'accuracy': 0.55}
step:  30 {'accuracy': 0.59}
step:  40 {'accuracy': 0.74}
step:  50 {'accuracy': 0.77}
step:  60 {'accuracy': 0.81}

Finish:  64 {'accuracy': 0.82}

可以看到，随着模型的训练，预测准确率从0.47上升到了0.82。

总结

本文使用Pytorch框架实现了Transformers库中的trainer模块对训练和验证的流程，自定义流程的好处在于灵活，但同样编写复杂；一般而言，在没有太多的模型内部更改时，可以优先考虑使用trainer，结合datasets库直接做到Huggingface全家桶的一站式开发，非常方便。