如何使用 Huggingface 的 Trainer 进行单轮训练？

Question

我想定义一个 Huggingface Trainer 对象，其中包含一组训练参数，包括在给定的一组历元上进行学习率退火的线性时间表，然后继续一次训练一个历元，以保持训练器的状态（优化器/计划/等..）历经了各个时期。

我想做的事情看起来像：

for i in range(epochs):
    data = modify_data()
    trainer.train_dataset = data["train"]
    trainer.train_one_epoch()

如果我只是在 TrainingArguments 中将

num_train_epochs

参数设置为 1，则学习率调度程序会将两个 epoch 之间的学习率设为 0.0，从而使训练在第一个 epoch 后毫无用处。 如果我只是在每次迭代时创建一个新的训练器，我就会失去学习率计划的状态。

有谁知道如何一次为一个时期创建一个 Trainer 对象训练？

PS：我正在使用 pytorch 后端

编辑以添加更多详细信息。

Answer 1

可以在TrainingArguments中添加训练epoc的数量，这里是一个例子：

from transformers import TrainingArguments, Trainer
training_args = TrainingArguments("test-trainer", num_train_epochs=1)

model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=2)

trainer = Trainer(
    model,
    training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)
trainer.train()

Answer 2

如果您想在一个时期后停止训练，您可以定义一个扩展 TrainerCallback 的类（请参阅code或查看huggingfaceimplements作为参考）。例如：

class StopCallback(TrainerCallback):
    def on_epoch_end(self, args, state, control, logs=None, **kwargs):
        control.should_training_stop = True

之后，您可以在定义 Trainer 时将此 StopCallback 作为参数传递。例如：

trainer = Trainer(
model,
training_args,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["validation"],
data_collator=data_collator,
tokenizer=tokenizer,
compute_metrics=compute_metrics,
callbacks=[StopCallback()]

）