携手创作,共同成长!这是我参与「掘金日新计划 · 8 月更文挑战」的第30天,简单的文本分类任务:使用Huggingface Trainer实现 为大家展示了如何使用Huggingface Transformers库来快速上手实现一个现代的文本分类模型的范式,虽然trainer很好用,但由于其封装的太过完善,所以在某些修改方面显得不是那么方便,我们有必要自己使用Pytorch手写一下整个流程。
任务简介
与前一篇博客相同,仍进行IMDB数据集上的情感分析。不过本次实验的不同之处在于:
- 为了快速实现,使用训练数据为1000条,测试为100条;
- 使用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())
整体的代码分为三个部分:
- 训练前的验证:用以检查模型在训练之前的性能表现。
- 模型训练:包括以下流程,
- 从Dataloader中读取数据
- 输入给模型进行前向计算
- 计算loss数值
- loss进行梯度回传
- 优化器进行参数更新
- 优化器梯度清零
- 训练完成后进行验证:检验模型的最终表现
得到的打印输出如下:
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全家桶的一站式开发,非常方便。