如何用HuggingFace Trainer API在TensorBoard同图展示训练与验证损失

核心思路 #

默认情况下，HuggingFace的TensorBoardCallback会为训练和验证分别创建独立的SummaryWriter，导致损失分属不同图表。通过共享同一个SummaryWriter实例，就能让训练损失（train/loss）和验证损失（eval/loss）被TensorBoard自动合并到同一张图中。

具体实现代码 #

假设你已有模型、数据集、Tokenizer等基础组件，只需添加以下关键部分：

1. 导入必要模块

from torch.utils.tensorboard import SummaryWriter
from transformers import TensorBoardCallback, Seq2SeqTrainingArguments, Seq2SeqTrainer

2. 创建共享的SummaryWriter
   指定统一的日志目录，建议和训练参数中的logging_dir保持一致：

log_dir = "./tb_logs"
tb_writer = SummaryWriter(log_dir=log_dir)

3. 初始化自定义TensorBoardCallback
   把共享的writer传入回调：

tb_callback = TensorBoardCallback(tb_writer=tb_writer)

4. 配置训练参数
   确保开启按步日志和验证，且日志步长与验证步长匹配，保证损失点对齐：

training_args = Seq2SeqTrainingArguments(
    output_dir="./results",
    logging_dir=log_dir,  # 与SummaryWriter路径一致
    logging_steps=100,  # 每100步记录训练损失
    evaluation_strategy="steps",  # 按步执行验证
    eval_steps=100,  # 每100步计算一次验证损失
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    # 其他你的自定义参数...
)

5. 初始化Trainer并传入回调

trainer = Seq2SeqTrainer(
    model=your_model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=your_tokenizer,
    callbacks=[tb_callback],  # 传入自定义回调
    # 其他你的自定义参数...
)

6. 训练后关闭writer

trainer.train()
tb_writer.close()

验证效果 #

启动TensorBoard（命令行输入tensorboard --logdir ./tb_logs），进入Scalars页面，就能看到train/loss和eval/loss被展示在同一张图表中，方便对比训练过程中的损失变化。

内容的提问来源于stack exchange，提问作者anna-kay