如果發現將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
。