Kaggle干货：自定义transformers数据集

transformers是现在NLP同学必备的库，但在使用的过程中主要的代码是需要自定义数据集，那么如何舒服的读取数据，并使用transformers进行训练模型呢？

本文的内容如下：

• 自定义文本数据集

• 使用transformers中trainer训练

• 使用Pytorch自定义训练流程

步骤1：读取IMDb Reviews

IMDb Reviews是比较常见的英文情感分类的数据集，主要完成文本多分类任务。我们首先需要读取数据，并处理数据。

• 下载数据

# 如下为shell命令
wget http://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz
tar -xf aclImdb_v1.tar.gz

• 划分数据集

from sklearn.model_selection import train_test_split
train_texts, val_texts, train_labels, val_labels = train_test_split(train_texts, train_labels, test_size=.2)

• 对文本进行编码

这里建议提前将文本进行编码，不建议动态进行编码，可以节约时间。

from transformers import DistilBertTokenizerFast
tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')

train_encodings = tokenizer(train_texts, truncation=True, padding=True)
val_encodings = tokenizer(val_texts, truncation=True, padding=True)
test_encodings = tokenizer(test_texts, truncation=True, padding=True)

• 自定义数据集

比较常见的自定义数据集的方法，需要修改__init__、__getitem__和__len__函数。

import torch

class IMDbDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

    def __len__(self):
        return len(self.labels)

train_dataset = IMDbDataset(train_encodings, train_labels)
val_dataset = IMDbDataset(val_encodings, val_labels)
test_dataset = IMDbDataset(test_encodings, test_labels)

步骤2：使用trainer训练

如果你的数据集定义为上述的格式，就可以非常方面的使用transformers中trainer训练。

from transformers import DistilBertForSequenceClassification, Trainer, TrainingArguments

# 训练细节
training_args = TrainingArguments(
    output_dir='./results',          # output directory
    num_train_epochs=3,              # total number of training epochs
    per_device_train_batch_size=16,  # batch size per device during training
    per_device_eval_batch_size=64,   # batch size for evaluation
    warmup_steps=500,                # number of warmup steps for learning rate scheduler
    weight_decay=0.01,               # strength of weight decay
    logging_dir='./logs',            # directory for storing logs
    logging_steps=10,
)

# 加载模型
model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")

trainer = Trainer(
    model=model,                         # the instantiated  Transformers model to be trained
    args=training_args,                  # training arguments, defined above
    train_dataset=train_dataset,         # training dataset
    eval_dataset=val_dataset             # evaluation dataset
)
trainer.train()

使用上述几行代码就可以自动帮助我们训练过程，这种方式有如下几个优点：

1. 速度比自己写正向/反向传播快；

2. 可以快速完成训练，保存调参；

3. 默认支持多卡训练；

当然也存在对应的缺点：

1. 自定义模型需要改写为官方的格式；

2. 默认没有对抗训练等其他操作；

步骤3：手动正向/反向传播

如果想要手动控制训练，则参考如下代码，需要自己完成损失计算、参数更新的过程。

from torch.utils.data import DataLoader
from transformers import DistilBertForSequenceClassification, AdamW

# 加载模型
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
model.to(device)
model.train()

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)

# 定义优化器
optim = AdamW(model.parameters(), lr=5e-5)

# epoch循环
for epoch in range(3):
    # batch 循环
    for batch in train_loader:
        optim.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs[0]
        loss.backward()
        optim.step()

model.eval()

手动实现训练细节更加可控，但需要确定的细节更多。所以如果你是初学者可以推荐使用trainer的方式。

步骤4：如何自定义模型？

在transformers中内置了常见的NLP任务，也对常见的模型进行了封装。比如：

{BERT名字}ForMaskedLM：MaskLM任务
{BERT名字}ForSequenceClassification：文本分类
{BERT名字}ForMultipleChoice：多项选择
{BERT名字}ForTokenClassification：实体分类
{BERT名字}ForQuestionAnswering：文本问答

当然如果你想要自定义模型，不使用已有的封装的模型，如何使用呢？可以参考已有的模型定义方式，如下：

• 步骤1：输入文本得到输出

• 步骤2：根据输入计算损失

• 步骤3：将输出和损失一起返回

outputs = self.roberta(
    input_ids,
    attention_mask=attention_mask,
    token_type_ids=token_type_ids,
)
sequence_output = outputs[0]
logits = self.classifier(sequence_output)

loss = None
if labels is not None:
    if self.num_labels == 1:
        #  We are doing regression
        loss_fct = MSELoss()
        loss = loss_fct(logits.view(-1), labels.view(-1))
    else:
        loss_fct = CrossEntropyLoss()
        loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

if not return_dict:
    output = (logits,) + outputs[2:]
    return ((loss,) + output) if loss is not None else output

return SequenceClassifierOutput(
    loss=loss,
    logits=logits,
    hidden_states=outputs.hidden_states,
    attentions=outputs.attentions,
)

参考资料

• https://huggingface.co/transformers/v3.2.0/custom_datasets.html

• https://huggingface.co/transformers/v3.2.0/_modules/transformers/modeling_roberta.html

往期精彩回顾




适合初学者入门人工智能的路线及资料下载(图文+视频)机器学习入门系列下载机器学习及深度学习笔记等资料打印《统计学习方法》的代码复现专辑机器学习交流qq群955171419，加入微信群请扫码