Trainer
Published:
Trainer学习笔记
Trainer是由huggingface开发的模型训练框架。接下来将分几个步骤介绍框架细节，主要内容参考官方博客：
Tokenizer and Dataset初始化
使用AutoTokenizer即可加载分词器：
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-cased") # 可以换成本地的地址
def tokenize_function(examples):
#
return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512) # 转换为input_ids和attention_mask
可以使用transformers预设的Trainer任务，例如BERT分类：
from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained("google-bert/bert-base-cased", num_labels=5)
HfArgumentParser
通过HfArgumentParser，可以将命令行的参数放到不同的数据类下，例如TrainingArguments, ModelArguments。 常见的模版：
from dataclasses import dataclass, field
from transformers import HfArgumentParser
@dataclass
class ModelArguments:
'''模型相关的参数
'''
model_name_or_path: Optional[str] = field(default="BertModel/")
tokenizer_name: Optional[str] = field(default="BertModel/")
load_plm: Optional[str] = field(default='CLModel/BertContrastive.aol')
load_model: Optional[str] = field(default='Ranking/ROUTE/PointBertSessionSearch.aol')
ltr: Literal["point", "pair"] = field(default="point")
@dataclass
class DataArguments:
'''数据相关的参数
'''
data_path: str = field(default="data/aol/", metadata={"help": "Path to the data."})
task: str = field(default="aol")
# training_args包含以下内容，可以通过argumentparser自动赋值
# 注意！trainer默认优化器为AdamW，以及使用线性衰减的学习率调度器！
'''
training_args = TrainingArguments(
output_dir="./results",
overwrite_output_dir=True,
num_train_epochs=3,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
warmup_steps=500,
weight_decay=0.01,
logging_dir="./logs",
logging_steps=10,
save_steps=10,
evaluation_strategy="steps",
eval_steps=500,
learning_rate=2e-5,
load_best_model_at_end=True
)
'''
def main():
parser = transformers.HfArgumentParser((TrainingArguments, ModelArguments))
# 从命令行解析参数，放到不同的数据类下
training_args, model_args = parser.parse_args_into_dataclasses()
Trainer
我们只需要定义Trainer就可以开始训练：
from transformers import Trainer
trainer = Trainer(
model,
training_args, # 只需要training_args，其他的args用于设定超参数
train_dataset=tokenized_datasets["train"], # 注意这里需要用Dataset.map方法预处理，例如通过tokenize_function将文本转换为input_ids等等
eval_dataset=tokenized_datasets["validation"],
data_collator=data_collator,# 缺省值为DataCollatorWithPadding
tokenizer=tokenizer,
callbacks=[swanlab_callback], # 回调函数，记录loss等
compute_metrics=compute_metrics
)
trainer.train()
train_dataset和eval_dataset可以通过自定义Dataset类来替代。 compute_metrics可以按如下形式定义：
import evaluate
def compute_metrics(eval_preds):
# eval_preds是一个EvalPrediction对象，包含两个属性：prediction, label_ids
metric = evaluate.load("glue", "mrpc")
logits, labels = eval_preds
predictions = np.argmax(logits, axis=-1)
return metric.compute(predictions=predictions, references=labels) # 返回accuracy & f1
metric.compute()也可以和trainer解耦，在每一步添加到缓存列表中，并在运行结束后再计算：
import evaluate
metric = evaluate.load("glue", "mrpc")
model.eval()
for batch in eval_dataloader:
batch = {k: v.to(device) for k, v in batch.items()}
with torch.no_grad():
outputs = model(**batch)
logits = outputs.logits
predictions = torch.argmax(logits, dim=-1)
metric.add_batch(predictions=predictions, references=batch["labels"])
metric.compute()
如果不使用evaluate类，也可以使用自定义的方式：
def compute_metrics(p: EvalPrediction):
# 获取预测结果和真实标签
preds = np.argmax(p.predictions, axis=1)
labels = p.label_ids
# 计算准确率和 F1 分数
accuracy = accuracy_score(labels, preds)
f1 = f1_score(labels, preds, average='weighted')
# 返回结果字典
return {
'accuracy': accuracy,
'f1': f1,
}
现在我们回到trainer传递的参数model。
我们通过一个具体任务加载模型：
from transformers import AutoModelForSequenceClassification
以此方式定义的模型都有期望的输入和输出。输入的key为input_ids, attention_mask, token_type_ids，而输出的key有labels。如果提供了输出的key，则模型在被调用时会返回loss:
outputs = model(**batch)
print(outputs.loss, outputs.logits.shape)