使用Bert模型优化Padding策略：加速文本分类训练

摘要

本文探讨了如何通过优化Padding策略，提高基于Bert的文本分类模型的训练速度。我们比较了两种不同的Padding方式：一种是将输入数据统一填充到最大长度512，另一种是只将每个Batch中的数据填充到Batch中最长的样本长度。

通过实验结果证明，后者显著减少了训练时间，且不影响模型的性能。

介绍

在本节中，我们研究 🤗 Transformers 中标记器的功能。分析 batch 数据 的padding。
下述 tokenizer 代码，会把每一条数据，padding 到最大长度 512。

tokenized_inputs = tokenizer(
    item["text"],
    max_length=512,
    padding="max_length",
    truncation=True,
)

下述代码仅仅进行 tokenize 化，只进行截断，但暂时还 不进行padding。把 padding 操作留到 后续的 batch 数据处理上，这样只需要padding 到 batch 数据里最长的数据长度，而无需 padding 到512。由于减少了input_ids的数据长度，所以在一定程度上可以加快模型的训练和推理速度。

tokenized_inputs = tokenizer(
    item["text"],
    max_length=512,
    truncation=True,
)

导包

import os
import evaluate
import numpy as np
from datasets import load_dataset
from transformers import (
    Trainer,
    TrainingArguments,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    DataCollatorWithPadding,
)
from dataclasses import dataclass

如果无法连接huggingface，利用c l a s h 的代理。

os.environ['HTTP_PROXY'] = 'http://127.0.0.1:7890'
os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:7890'

模型

使用 Bert 模型，进行实验。

model_name = "bert-base-chinese"

bert = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    trust_remote_code=True,
    num_labels=2,
)

tokenizer = AutoTokenizer.from_pretrained(model_name)

dataset

使用 huggingface 平台的 中文二分类数据集。

ds = load_dataset("lansinuote/ChnSentiCorp")

使用 transformers 的 DataCollatorWithPadding。

data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

def tokenize_func_pad(item):
    global tokenizer
    tokenized_inputs = tokenizer(
        item["text"],
        max_length=512,
        truncation=True,
        padding="max_length"
    )
    return tokenized_inputs

batch_dataset_pad = ds["train"].map(
    tokenize_func_pad, remove_columns=["text"]
)
batch_dataset_pad.set_format("torch")
batch_dataset_pad["input_ids"].shape

输出结果:

torch.Size([9600, 512])

batch_dataset_pad["input_ids"].shape 是 torch.Size([9600, 512])，每一条数据都padding 到512长度。这里面会包含大量的<pad>。

只需要在每一个 batch 数据中，把其中所有的样本 padding 到这一个batch最长的样本长度即可。无需把所有的样本都 padding 到 512。

def tokenize_func(item):
    global tokenizer
    tokenized_inputs = tokenizer(
        item["text"],
        max_length=512,
        truncation=True,
    )
    return tokenized_inputs

batch_dataset = ds["train"].map(
	tokenize_func, 
	remove_columns=["text"]
)

data = [batch_dataset[i] for i in range(16)]
data_collator(data)["input_ids"].shape

输出结果：

torch.Size([16, 126])

上述结果展示了，把当前batch里的数据，padding到当前batch的最大长度126。

data_collator 可传递给 DataLoader 和 Trainer。

BertCLS

下述是 Bert 封装好的一个工具类。方便读者使用，训练模型、评估、预测。

@dataclass
class BertCLS:
    def __init__(
        self,
        model,
        tokenizer,
        train_dataset=None,
        eval_dataset=None,
        output_dir="output",
        epoch=3,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset

        self.data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

        self.args = self.get_args(output_dir, epoch)

        self.trainer = Trainer(
            model=self.model,
            args=self.args,
            train_dataset=self.train_dataset,
            eval_dataset=self.eval_dataset,
            data_collator=self.data_collator,
            # compute_metrics=compute_metrics,
            tokenizer=tokenizer,
        )

    def get_args(self, output_dir, epoch):
        if self.eval_dataset:
            args = TrainingArguments(
                output_dir=output_dir,
                evaluation_strategy="epoch",
                save_strategy="epoch",
                save_total_limit=3,
                learning_rate=2e-5,
                num_train_epochs=epoch,
                weight_decay=0.01,
                per_device_train_batch_size=32,
                per_device_eval_batch_size=16,
                # logging_steps=16,
                save_safetensors=True,
                overwrite_output_dir=True,
                load_best_model_at_end=True,
            )
        else:
            args = TrainingArguments(
                output_dir=output_dir,
                evaluation_strategy="no",
                save_strategy="epoch",
                save_total_limit=3,
                learning_rate=2e-5,
                num_train_epochs=epoch,
                weight_decay=0.01,
                per_device_train_batch_size=32,
                per_device_eval_batch_size=16,
                # logging_steps=16,
                save_safetensors=True,
                overwrite_output_dir=True,
                # load_best_model_at_end=True,
            )
        return args

    def set_args(self, args):
        """
        从外部重新设置 TrainingArguments，args 更新后，trainer也进行更新
        """
        self.args = args

        self.trainer = Trainer(
            model=self.model,
            args=self.args,
            train_dataset=self.train_dataset,
            eval_dataset=self.eval_dataset,
            data_collator=self.data_collator,
            # compute_metrics=compute_metrics,
            tokenizer=self.tokenizer,
        )

    def train(self, epoch=None, over_write=False):
        if epoch:
            self.args.num_train_epochs = epoch
            
        best_model_path = os.path.join(self.args.output_dir, "best_model")

        if over_write or not os.path.exists(best_model_path):
            self.trainer.train()
            self.trainer.save_model(best_model_path)
        else:
            print(
                f"预训练权重 {best_model_path} 已存在，且over_write={over_write}。不启动模型训练！"
            )

    def eval(self, eval_dataset):
        predictions = self.trainer.predict(eval_dataset)
        preds = np.argmax(predictions.predictions, axis=-1)
        metric = evaluate.load("glue", "mrpc")
        return metric.compute(predictions=preds, references=predictions.label_ids)

    def pred(self, pred_dataset):
        predictions = self.trainer.predict(pred_dataset)
        preds = np.argmax(predictions.predictions, axis=-1)
        return pred_dataset.add_column("pred", preds)

速度对比

只给 BertCLS 训练数据集，开始训练模型。如果要看模型训练的评估结果，输入评估数据集即可。

padding 到最大长度：

BertCLS(bert, tokenizer, batch_dataset_pad, epoch=1).train()

 [300/300 03:09, Epoch 1/1]

padding 到每个batch的最大长度：

BertCLS(bert, tokenizer, batch_dataset, epoch=1).train()

 [300/300 02:19, Epoch 1/1]

padding 到 512个长度，训练3分9秒结束。如果只padding到每个batch的最大长度，训练2分19秒结束。

代码开源地址

https://github.com/JieShenAI/csdn/blob/main/24/09/tokenizer_pad/pad_vs.ipynb