如果发现将SummaryWriter
直接添加到callbacks
不被支持,您可以采取另一种方式来集成TensorBoard记录功能,即通过自定义回调函数来实现。Hugging Face Transformers库允许用户自定义训练回调,这可以用来在训练过程中向TensorBoard写入日志。
下面是如何创建一个简单的自定义回调以与TensorBoard一起使用的示例:
首先,确保你有SummaryWriter
的实例在训练开始前被初始化,并且在训练结束后被关闭。然后,定义一个自定义回调类来集成TensorBoard日志记录逻辑:
from transformers import TrainerCallback
from torch.utils.tensorboard import SummaryWriterclass TensorBoardCallback(TrainerCallback):def __init__(self, writer):self.writer = writerdef on_log(self, args, state, control, logs=None, **kwargs):if logs is not None:for k, v in logs.items():if isinstance(v, (int, float)):self.writer.add_scalar(k, v, state.global_step)self.writer.flush()# 初始化SummaryWriter
tb_writer = SummaryWriter(log_dir=training_args.logging_dir)# 自定义回调实例化
custom_callback = TensorBoardCallback(writer=tb_writer)# 定义训练器,将自定义回调添加到callbacks列表中
trainer = Trainer(model=model, # 待训练模型args=training_args, # 训练参数data_collator=collater, # 数据校准器train_dataset=tokenized_train_dataset,# 训练集eval_dataset=tokenized_valid_dataset, # 验证集compute_metrics=compute_metric, # 计算自定义指标callbacks=[custom_callback], # 添加自定义TensorBoard回调
)# 开始训练
trainer.train()# 训练结束后关闭TensorBoard SummaryWriter
tb_writer.close()
在这个示例中,我们创建了一个名为TensorBoardCallback
的类,该类继承自TrainerCallback
。在on_log
方法中,我们将训练日志中的每个标量值写入到TensorBoard中,这样您就可以在TensorBoard中可视化这些指标了。记得在训练结束后关闭SummaryWriter
。