之前只闻 transformers 超厉害超好用，但是没有实际用过。之前涉及到 bert 类模型都是直接手写或是在别人的基础上修改。但这次由于某些原因，需要快速训练一个简单的文本分类模型。其实这种场景应该挺多的，例如简单的 POC 或是临时测试某些模型。

我的需求很简单：用我们自己的数据集，快速训练一个文本分类模型，验证想法。

我觉得如此简单的一个需求，应该有模板代码。但实际去搜的时候发现，官方文档什么时候变得这么多这么庞大了？还多了个 Trainer API？瞬间让我想起了 Pytorch Lightning 那个坑人的同名 API。但可能是时间原因，找了一圈没找到适用于自定义数据集的代码，都是用的官方、预定义的数据集。

所以弄完后，我决定简单写一个文章，来说下这原本应该极其容易解决的事情。

数据

假设我们数据的格式如下：

1
2
3

0 第一个句子
1 第二个句子
0 第三个句子

即每一行都是 label sentence 的格式，中间空格分隔。并且我们已将数据集分成了 train.txt 和 val.txt 。

代码

加载数据集

首先使用 datasets 加载数据集：

1 2	from datasets import load_dataset dataset = load_dataset('text', data_files={'train': 'data/train_20w.txt', 'test': 'data/val_2w.txt'})

加载后的 dataset 是一个 DatasetDict 对象：

DatasetDict({
    train: Dataset({
        features: ['text'],
        num_rows: 3
    })
    test: Dataset({
        features: ['text'],
        num_rows: 3
    })
})

类似 tf.data ，此后我们需要对其进行 map ，对每一个句子进行 tokenize、padding、batch、shuffle：

def tokenize_function(examples):
    labels = []
    texts = []
    for example in examples['text']:
        split = example.split(' ', maxsplit=1)
        labels.append(int(split[0]))
        texts.append(split[1])
    tokenized = tokenizer(texts, padding='max_length', truncation=True, max_length=32)
    tokenized['labels'] = labels
    return tokenized

tokenized_datasets = dataset.map(tokenize_function, batched=True)
train_dataset = tokenized_datasets["train"].shuffle(seed=42)
eval_dataset = tokenized_datasets["test"].shuffle(seed=42)

根据数据集格式不同，我们可以在 tokenize_function 中随意自定义处理过程，以得到 text 和 labels。注意 batch_size 和 max_length 也是在此处指定。处理完我们便得到了可以输入给模型的训练集和测试集。

训练

model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2, cache_dir='data/pretrained')
training_args = TrainingArguments('ckpts', per_device_train_batch_size=256, num_train_epochs=5)
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset
)
trainer.train()

你可以根据情况修改训练 batchsize per_device_train_batch_size 。

增加准确率显示

我们在训练的时候一般会监测测试准确率来评估模型性能，而 transformers 在训练过程中默认是不会输出准确率的，而且训练完也不会输出的。这样的话我们想要一个准确率的话，只能再手动加载一下模型然后走一下预测，略显麻烦。

但 transformers 也是支持计算并输出准确率的，我们可以为 Trainer 指定 compute_metrics 参数。

compute_metrics 参数必须是一个函数，用于计算准确率等 metrics 的函数。该函数的输入是 transformers.EvalPrediction 对象，包含模型的输出（logits）和正确标签，其本质上是一个 namedtuple，相应的 field 为 predictions 和 label_ids；输出必须是一个字典，key 为 metric name，value 为 metric value。

关于 metric 的计算，datasets 实际上已经为我们提供了一些内置函数。你可以用 datasets.list_metrics() 来获取目前所有可用的 metric。但是在 load_metric() 时，需要从 GitHub 下载处理程序，鉴于国内网络状况，这步通常都会卡住：

1
2

# https://github.com/huggingface/datasets/blob/21bfd0d3f5ff3fbfd691600e2c7071a167816cdf/src/datasets/config.py#L21
REPO_METRICS_URL = "https://raw.githubusercontent.com/huggingface/datasets/{revision}/metrics/{path}/{name}"

解决这种情况有几种办法：

挂梯子。
load_metric() 支持从本地加载计算程序，所以你可以把 metric 计算代码放到你本地，然后把地址传进去。
不使用 load_metric()，而是我们自己根据 predictions 和 label_ids 来计算。

本文接下来就是使用最后一种方法，较为灵活。我们可以使用 scikit-learn 来计算这些 metric。实际上 datasets 中的 accuracy 也是使用 sklearn.metrics.accuracy_score 来计算的。

来看代码：

from sklearn.metrics import accuracy_score, f1_score


def compute_metrics(eval_pred) -> dict:
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    acc = accuracy_score(labels, predictions)
    f1 = f1_score(labels, predictions, average='micro')
    return {"accuracy": acc, 'f1': f1}


training_args = TrainingArguments(
    # 其他参数
    evaluation_strategy="epoch",
    # 其他参数
)

trainer = Trainer(
    # 其他参数
    args=training_args,
    eval_dataset=eval_dataset,
    compute_metrics=compute_metrics,  # <-- 计算metric
    # 其他参数
)

注意一定要加上上面 TrainingArguments 中的 evaluation_strategy="epoch"，该参数默认是 "no"，即不进行 evaluation。我们此处指定为 "epoch" 表示在每个 epoch 结束时进行 evaluation。其他可选的值为 "steps"，表示每 eval_steps 进行一次 evaluation，默认为 500 steps。

然后运行我们即可看到类似如下的输出：

***** Running training *****
  Num examples = 43410
  Num Epochs = 5
  Instantaneous batch size per device = 48
  Total train batch size (w. parallel, distributed & accumulation) = 48
  Gradient Accumulation steps = 1
  Total optimization steps = 4525
{'loss': 0.82, 'learning_rate': 4.447513812154696e-05, 'epoch': 0.55}
 20%|████                           | 905/4525 [06:20<21:36,  2.79it/s]
 The following columns in the evaluation set  don't have a corresponding argument in `BertForSequenceClassification.forward` and have been ignored: text.        
***** Running Evaluation *****
  Num examples = 5426
  Batch size = 8
{'eval_loss': 0.7136639952659607, 'eval_accuracy': 0.7051234795429414, 'eval_f1': 0.7051234795429414, 'eval_runtime': 20.5673, 'eval_samples_per_second': 263.817, 'eval_steps_per_second': 33.014, 'epoch': 1.0}

我们可以看到在第一个 epoch 结束之后进行了 evaluation，accuracy 和 f1 也被正确返回了（会加上 eval_ 前缀）。

只保存性能最好的 checkpoint

根据 save_strategy 的不同，训练时默认每隔一定时间段就保存一次模型 checkpoint。如果训练 epochs 比较多，会保存很多 ckpt。但有时我们硬盘空间有限，或者由于其他原因不想保存这么多的 ckpt，只想保存最佳模型的。

transformers 也可以很方便地实现这个功能。

严格来说会保存两个 ckpt：一个最佳的，一个最后的（用于接续训练）。

默认情况下，我们只需要给 TrainingArguments 多加两个参数：

load_best_model_at_end=True：训练结束加载最佳模型。
save_total_limit=1：总共保存 1 个模型 ckpt（实际是两个）。

那么如何判断最佳呢？

通过 metric_for_best_model 和 greater_is_better 来共同判断。要想判断最佳，我们首先需要知道评判标准是什么，这就是前者的作用。默认是 loss，在 eval_dataset 上的 loss，你也可以指定为 compute_metrics() 所返回的 metric name（带不带 eval_ 都行）。其次我们需要知道这个标准是越大越好还是越小越好，这就是后者的作用。如果是标准是 loss，那么会自动设置为 False，因为 loss 是越小越好。但如果你指定为其他的标准，记得手动设置下这个参数。

来看代码：

training_args = TrainingArguments(
    # 其他参数
    load_best_model_at_end=True,
    save_total_limit=1,
    # 其他参数
)

完整代码

完整代码见 GitHub。

使用 Transformers 在你自己的数据集上训练文本分类模型

目录

背景

数据

代码

加载数据集

训练

增加准确率显示

只保存性能最好的 checkpoint

完整代码

END

Alan Lee