摘要
本文探讨了如何通过优化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