摘要

本文探讨了如何通过优化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 tokenizertokenized_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 tokenizertokenized_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 = modelself.tokenizer = tokenizerself.train_dataset = train_datasetself.eval_dataset = eval_datasetself.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 argsdef set_args(self, args):"""从外部重新设置 TrainingArguments，args 更新后，trainer也进行更新"""self.args = argsself.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 = epochbest_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